Update app.py

app.py

@@ -115,7 +115,7 @@ class GPUSatelliteModelGenerator:
         road_mask = cp.zeros((height * width,), dtype=bool)
         water_mask = cp.zeros((height * width,), dtype=bool)
         
-        # Vectorized color matching
+        # Vectorized color matching for predefined colors
         for ref_hsv in self.shadow_colors_hsv:
             shadow_mask |= self.gpu_color_distance_hsv(hsv_pixels.T, ref_hsv, self.shadow_tolerance)
             
@@ -125,25 +125,54 @@ class GPUSatelliteModelGenerator:
         for ref_hsv in self.water_colors_hsv:
             water_mask |= self.gpu_color_distance_hsv(hsv_pixels.T, ref_hsv, self.water_tolerance)
         
-        # Apply masks
-        output_flat = output.reshape(-1, 3)
-        output_flat[shadow_mask] = self.colors['black']
-        output_flat[water_mask] = self.colors['blue']
-        output_flat[road_mask] = self.colors['gray']
-        
-        # Vegetation and building detection
+        # Convert HSV to normalized values for vegetation and terrain detection
         h, s, v = hsv_pixels.T
         h = h * 2  # Convert to 0-360 range
         s = s / 255
         v = v / 255
         
+        # Create masks for vegetation and terrain
         vegetation_mask = (h >= 40) & (h <= 150) & (s >= 0.15)
-        building_mask = ~(shadow_mask | water_mask | road_mask | vegetation_mask)
+        terrain_mask = (h >= 10) & (h <= 30) & (s >= 0.15)
         
+        # Building mask (everything that's not another category)
+        building_mask = ~(shadow_mask | water_mask | road_mask | vegetation_mask | terrain_mask)
+        
+        # Apply all masks to create final output
+        output_flat = output.reshape(-1, 3)
+        output_flat[shadow_mask] = self.colors['black']
+        output_flat[water_mask] = self.colors['blue']
+        output_flat[road_mask] = self.colors['gray']
         output_flat[vegetation_mask] = self.colors['green']
+        output_flat[terrain_mask] = self.colors['brown']
         output_flat[building_mask] = self.colors['white']
         
-        return output.reshape(height, width, 3)
+        # Reshape back to image dimensions
+        segmented = output.reshape(height, width, 3)
+        
+        # Handle isolated building pixels using GPU operations
+        kernel = cp.ones((3, 3), dtype=bool)
+        kernel[1, 1] = False
+        
+        # For each color except white (buildings)
+        for color_name, color_value in self.colors.items():
+            if cp.array_equal(color_value, self.colors['white']):
+                continue
+                
+            # Create mask for current color
+            color_mask = cp.all(segmented == color_value, axis=2)
+            
+            # Dilate the mask
+            dilated = binary_dilation(color_mask, structure=kernel)
+            
+            # Find building pixels that are mostly surrounded by this color
+            building_pixels = cp.all(segmented == self.colors['white'], axis=2)
+            surrounded = dilated & building_pixels
+            
+            # Update those pixels to the current color
+            segmented[surrounded] = color_value
+        
+        return segmented
         
     def estimate_heights_gpu(self, img, segmented):
         """GPU-accelerated height estimation"""