Update app.py

app.py

@@ -82,24 +82,30 @@ class GPUSatelliteModelGenerator:
 
     @staticmethod
     def gpu_color_distance_hsv(pixel_hsv, reference_hsv, tolerance):
-        """GPU-accelerated HSV color distance calculation"""
+        """HSV color distance calculation"""
         pixel_h = pixel_hsv[0] * 2
         pixel_s = pixel_hsv[1] / 255
         pixel_v = pixel_hsv[2] / 255
         
+        # Calculate circular hue difference
         hue_diff = cp.minimum(cp.abs(pixel_h - reference_hsv[0]),
-                            360 - cp.abs(pixel_h - reference_hsv[0]))
+                             360 - cp.abs(pixel_h - reference_hsv[0]))
+        
+        # Calculate saturation and value differences with weighted importance
         sat_diff = cp.abs(pixel_s - reference_hsv[1])
         val_diff = cp.abs(pixel_v - reference_hsv[2])
         
+        # Combined distance check with adjusted weights
         return cp.logical_and(
-            cp.logical_and(hue_diff <= tolerance['hue'],
-                          sat_diff <= tolerance['sat']),
+            cp.logical_and(
+                hue_diff <= tolerance['hue'],
+                sat_diff <= tolerance['sat']
+            ),
             val_diff <= tolerance['val']
         )
 
     def segment_image_gpu(self, img):
-        """GPU-accelerated image segmentation"""
+        """GPU-accelerated image segmentation with improved road and shadow detection"""
         # Transfer image to GPU
         gpu_img = cp.asarray(img)
         gpu_hsv = cp.asarray(cv2.cvtColor(img, cv2.COLOR_BGR2HSV))
@@ -107,38 +113,65 @@ class GPUSatelliteModelGenerator:
         height, width = img.shape[:2]
         output = cp.zeros_like(gpu_img)
         
-        # Vectorized color matching on GPU
+        # Create a sliding window view for neighborhood analysis
+        pad = 2  # Equivalent to window_size=5 in segment.py
+        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)
         
-        # Create masks for each category
+        # Initialize masks
         shadow_mask = cp.zeros((height * width,), dtype=bool)
         road_mask = cp.zeros((height * width,), dtype=bool)
         water_mask = cp.zeros((height * width,), dtype=bool)
         
-        # Vectorized color matching for predefined colors
+        # Improved color matching with adjusted tolerances
         for ref_hsv in self.shadow_colors_hsv:
-            shadow_mask |= self.gpu_color_distance_hsv(hsv_pixels.T, ref_hsv, self.shadow_tolerance)
-            
+            # Lower the threshold for shadows to catch more subtle variations
+            temp_tolerance = {
+                'hue': self.shadow_tolerance['hue'] * 1.2,  # Slightly increased tolerance
+                '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:
-            road_mask |= self.gpu_color_distance_hsv(hsv_pixels.T, ref_hsv, self.road_tolerance)
-            
+            # 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
+            }
+            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)
         
-        # Convert HSV to normalized values for vegetation and terrain detection
+        # Normalize HSV 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)
-        terrain_mask = (h >= 10) & (h <= 30) & (s >= 0.15)
+        # 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
+        
+        # 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)
+        
+        # Enhanced shadow detection using value component
+        dark_mask = (v <= 0.3)  # Detect very dark areas
+        shadow_mask |= dark_mask & ~(water_mask | road_mask)
         
         # 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
+        # Apply masks to create output
         output_flat = output.reshape(-1, 3)
         output_flat[shadow_mask] = self.colors['black']
         output_flat[water_mask] = self.colors['blue']
@@ -147,30 +180,31 @@ class GPUSatelliteModelGenerator:
         output_flat[terrain_mask] = self.colors['brown']
         output_flat[building_mask] = self.colors['white']
         
-        # Reshape back to image dimensions
         segmented = output.reshape(height, width, 3)
         
-        # Handle isolated building pixels using GPU operations
+        # Enhanced isolated pixel cleanup using morphological 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
+        # Two-pass cleanup for better results
+        for _ in range(2):
+            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
+                # 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
+                
+                # Update isolated pixels
+                segmented[surrounded] = color_value
         
         return segmented
         
@@ -324,8 +358,8 @@ def generate_and_process_map(prompt: str) -> tuple[str | None, np.ndarray | None
 
 # Create Gradio interface
 with gr.Blocks() as demo:
-    gr.Markdown("# GPU-Accelerated Text to Map")
-    gr.Markdown("Generate 3D maps and segmentation maps from text descriptions using FLUX and GPU-accelerated processing.")
+    gr.Markdown("# Text to Map")
+    gr.Markdown("Generate a 3D map from text!")
     
     with gr.Row():
         prompt_input = gr.Text(