Update app.py

app.py

@@ -39,6 +39,40 @@ class GPUSatelliteModelGenerator:
             [47, 94, 100],
             [73, 131, 135]
         ])
+
+        # Output colors (BGR for OpenCV)
+
+        self.roof_colors = cp.array([
+            [191, 148, 124],
+            [190, 142, 121],
+            [184, 154, 139],
+            [178, 118, 118],
+            [164, 109, 107],
+            [155, 113, 105],
+            [153, 111, 106],
+            [155, 95, 96],
+            [135, 82, 87],
+            [117, 82, 78],
+            [113, 62, 50],
+            [166, 144, 135]
+        ])
+        
+        # Convert roof colors to HSV
+        self.roof_colors_hsv = cp.asarray(cv2.cvtColor(
+            self.roof_colors.get().reshape(-1, 1, 3).astype(np.uint8),
+            cv2.COLOR_RGB2HSV
+        ).reshape(-1, 3))
+        
+        # Normalize roof HSV values
+        self.roof_colors_hsv[:, 0] = self.roof_colors_hsv[:, 0] * 2
+        self.roof_colors_hsv[:, 1:] = self.roof_colors_hsv[:, 1:] / 255
+        
+        # Add roof tolerance (tighter than terrain to avoid confusion)
+        self.roof_tolerance = {
+            'hue': 8,  # Tighter hue tolerance to differentiate from terrain
+            'sat': 0.15,
+            'val': 0.15
+        }
         
         # Convert reference colors to HSV on GPU
         self.shadow_colors_hsv = cp.asarray(cv2.cvtColor(
@@ -66,14 +100,15 @@ class GPUSatelliteModelGenerator:
         self.road_tolerance = {'hue': 10, 'sat': 0.12, 'val': 0.15}
         self.water_tolerance = {'hue': 20, 'sat': 0.15, 'val': 0.20}
         
-        # Output colors (BGR for OpenCV)
+        # Colors dictionary in [B, G, R]
         self.colors = {
             'black': cp.array([0, 0, 0]),     # Shadows
             'blue': cp.array([255, 0, 0]),    # Water
             'green': cp.array([0, 255, 0]),   # Vegetation
             'gray': cp.array([128, 128, 128]), # Roads
             'brown': cp.array([0, 140, 255]),  # Terrain
-            'white': cp.array([255, 255, 255]) # Buildings
+            'white': cp.array([255, 255, 255]), # Buildings
+            'salmon': cp.array([128, 128, 255])  # Roofs
         }
         
         self.min_area_for_clustering = 1000
@@ -105,7 +140,7 @@ class GPUSatelliteModelGenerator:
         )
 
     def segment_image_gpu(self, img):
-        """GPU-accelerated image segmentation with improved road and shadow detection"""
+        """GPU-accelerated image segmentation with roof detection"""
         # Transfer image to GPU
         gpu_img = cp.asarray(img)
         gpu_hsv = cp.asarray(cv2.cvtColor(img, cv2.COLOR_BGR2HSV))
@@ -114,40 +149,43 @@ class GPUSatelliteModelGenerator:
         output = cp.zeros_like(gpu_img)
         
         # Create a sliding window view for neighborhood analysis
-        pad = 2  # Equivalent to window_size=5 in segment.py
+        pad = 2
         gpu_hsv_pad = cp.pad(gpu_hsv, ((pad, pad), (pad, pad), (0, 0)), mode='edge')
         
         # Prepare flattened HSV data
         hsv_pixels = gpu_hsv.reshape(-1, 3)
         
-        # Initialize masks
+        # Initialize masks including roofs
         shadow_mask = cp.zeros((height * width,), dtype=bool)
         road_mask = cp.zeros((height * width,), dtype=bool)
         water_mask = cp.zeros((height * width,), dtype=bool)
+        roof_mask = cp.zeros((height * width,), dtype=bool)
         
-        # Improved color matching with adjusted tolerances
+        # Color matching for predefined categories
         for ref_hsv in self.shadow_colors_hsv:
-            # Lower the threshold for shadows to catch more subtle variations
             temp_tolerance = {
-                'hue': self.shadow_tolerance['hue'] * 1.2,  # Slightly increased tolerance
+                'hue': self.shadow_tolerance['hue'] * 1.2,
                 'sat': self.shadow_tolerance['sat'] * 1.1,
                 'val': self.shadow_tolerance['val'] * 1.2
             }
             shadow_mask |= self.gpu_color_distance_hsv(hsv_pixels.T, ref_hsv, temp_tolerance)
         
         for ref_hsv in self.road_colors_hsv:
-            # Adjusted road detection with focus on value component
             temp_tolerance = {
-                'hue': self.road_tolerance['hue'] * 1.3,  # Increased hue tolerance
-                'sat': self.road_tolerance['sat'] * 1.2,  # Increased saturation tolerance
-                'val': self.road_tolerance['val']  # Keep original value tolerance
+                'hue': self.road_tolerance['hue'] * 1.3,
+                'sat': self.road_tolerance['sat'] * 1.2,
+                'val': self.road_tolerance['val']
             }
             road_mask |= self.gpu_color_distance_hsv(hsv_pixels.T, ref_hsv, temp_tolerance)
         
         for ref_hsv in self.water_colors_hsv:
             water_mask |= self.gpu_color_distance_hsv(hsv_pixels.T, ref_hsv, self.water_tolerance)
         
-        # Normalize HSV values for vegetation and terrain detection
+        # Roof detection with specific color matching
+        for ref_hsv in self.roof_colors_hsv:
+            roof_mask |= self.gpu_color_distance_hsv(hsv_pixels.T, ref_hsv, self.roof_tolerance)
+        
+        # Normalize HSV values
         h, s, v = hsv_pixels.T
         h = h * 2  # Convert to 0-360 range
         s = s / 255
@@ -156,20 +194,22 @@ class GPUSatelliteModelGenerator:
         # Enhanced vegetation detection
         vegetation_mask = ((h >= 40) & (h <= 150) & (s >= 0.15))
         
-        # Enhanced terrain detection
-        terrain_mask = ((h >= 10) & (h <= 30) & (s >= 0.15)) | \
-                      ((h >= 25) & (h <= 40) & (s >= 0.1) & (v <= 0.8))  # Added brown-gray detection
+        # Refined terrain detection to avoid roof confusion
+        terrain_mask = (
+            ((h >= 15) & (h <= 35) & (s >= 0.15) & (s <= 0.6)) |  # Main terrain colors
+            ((h >= 25) & (h <= 40) & (s >= 0.1) & (v >= 0.5))     # Lighter terrain
+        ) & ~roof_mask  # Explicitly exclude roof areas
         
         # Apply brightness-based corrections for roads
-        gray_mask = (s <= 0.2) & (v >= 0.4) & (v <= 0.85)  # Detect grayish areas
-        road_mask |= gray_mask & ~(shadow_mask | water_mask | vegetation_mask | terrain_mask)
+        gray_mask = (s <= 0.2) & (v >= 0.4) & (v <= 0.85)
+        road_mask |= gray_mask & ~(shadow_mask | water_mask | vegetation_mask | terrain_mask | roof_mask)
         
-        # Enhanced shadow detection using value component
-        dark_mask = (v <= 0.3)  # Detect very dark areas
-        shadow_mask |= dark_mask & ~(water_mask | road_mask)
+        # Enhanced shadow detection
+        dark_mask = (v <= 0.3)
+        shadow_mask |= dark_mask & ~(water_mask | road_mask | roof_mask)
         
         # Building mask (everything that's not another category)
-        building_mask = ~(shadow_mask | water_mask | road_mask | vegetation_mask | terrain_mask)
+        building_mask = ~(shadow_mask | water_mask | road_mask | vegetation_mask | terrain_mask | roof_mask)
         
         # Apply masks to create output
         output_flat = output.reshape(-1, 3)
@@ -178,52 +218,56 @@ class GPUSatelliteModelGenerator:
         output_flat[road_mask] = self.colors['gray']
         output_flat[vegetation_mask] = self.colors['green']
         output_flat[terrain_mask] = self.colors['brown']
+        output_flat[roof_mask] = self.colors['salmon']
         output_flat[building_mask] = self.colors['white']
         
         segmented = output.reshape(height, width, 3)
         
-        # Enhanced isolated pixel cleanup using morphological operations
+        # Enhanced cleanup with roof consideration
         kernel = cp.ones((3, 3), dtype=bool)
         kernel[1, 1] = False
         
-        # Two-pass cleanup for better results
+        # Two-pass cleanup
         for _ in range(2):
             for color_name, color_value in self.colors.items():
                 if cp.array_equal(color_value, self.colors['white']):
                     continue
                 
-                # Create and dilate mask for current color
                 color_mask = cp.all(segmented == color_value, axis=2)
                 dilated = binary_dilation(color_mask, structure=kernel)
                 
-                # Find isolated building pixels
                 building_pixels = cp.all(segmented == self.colors['white'], axis=2)
                 neighbor_count = binary_dilation(color_mask, structure=kernel).astype(int)
                 
-                # More aggressive cleanup for truly isolated pixels
-                surrounded = (neighbor_count >= 5) & building_pixels  # At least 5 neighbors of same color
+                # Special handling for roofs - they should be more granular
+                if cp.array_equal(color_value, self.colors['salmon']):
+                    surrounded = (neighbor_count >= 4) & building_pixels  # Less aggressive for roofs
+                else:
+                    surrounded = (neighbor_count >= 5) & building_pixels
                 
-                # Update isolated pixels
                 segmented[surrounded] = color_value
         
         return segmented
-        
+            
     def estimate_heights_gpu(self, img, segmented):
-        """GPU-accelerated height estimation"""
+        """GPU-accelerated height estimation with roof consideration"""
         gpu_segmented = cp.asarray(segmented)
-        buildings_mask = cp.all(gpu_segmented == self.colors['white'], axis=2)
+        buildings_mask = cp.logical_or(
+            cp.all(gpu_segmented == self.colors['white'], axis=2),
+            cp.all(gpu_segmented == self.colors['salmon'], axis=2)
+        )
         shadows_mask = cp.all(gpu_segmented == self.colors['black'], axis=2)
         
         # Connected components labeling on GPU
         labeled_array, num_features = cp_label(buildings_mask)
         
         # Calculate areas using GPU
-        areas = cp.bincount(labeled_array.ravel())[1:]  # Skip background
+        areas = cp.bincount(labeled_array.ravel())[1:]
         max_area = cp.max(areas) if len(areas) > 0 else 1
         
         height_map = cp.zeros_like(labeled_array, dtype=cp.float32)
         
-        # Process each building
+        # Process each building/roof
         for label in range(1, num_features + 1):
             building_mask = (labeled_array == label)
             if not cp.any(building_mask):
@@ -232,12 +276,18 @@ class GPUSatelliteModelGenerator:
             area = areas[label-1]
             size_factor = 0.3 + 0.7 * (area / max_area)
             
+            # Check if this is a roof (salmon color)
+            is_roof = cp.any(cp.all(gpu_segmented[building_mask] == self.colors['salmon'], axis=1))
+            
+            # Adjust height for roofs (typically smaller residential buildings)
+            if is_roof:
+                size_factor *= 0.8  # Slightly lower height for residential buildings
+            
             # Calculate shadow influence
             dilated = binary_dilation(building_mask, structure=cp.ones((5,5)))
             shadow_ratio = cp.sum(dilated & shadows_mask) / cp.sum(dilated)
             shadow_factor = 0.2 + 0.8 * shadow_ratio
             
-            # Height calculation based on size and shadows
             final_height = size_factor * shadow_factor
             height_map[building_mask] = final_height