Update app.py

app.py

@@ -13,241 +13,214 @@ from sklearn.cluster import DBSCAN
 import trimesh
 
 class GPUSatelliteModelGenerator:
+    
     def __init__(self, building_height=0.05):
         self.building_height = building_height
         
-        # Move color arrays to GPU using cupy
-        self.shadow_colors = cp.array([
-            [31, 42, 76],
-            [58, 64, 92],
-            [15, 27, 56],
-            [21, 22, 50],
-            [76, 81, 99]
+        # Add grass and tree colors
+        self.grass_colors = cp.array([
+            [47, 70, 69],  # Light green grass
+            [40, 60, 55],
+            [45, 65, 60],
+            [50, 75, 65]
         ])
         
-        self.road_colors = cp.array([
-            [187, 182, 175],
-            [138, 138, 138], 
-            [142, 142, 129],
-            [202, 199, 189]
+        self.tree_colors = cp.array([
+            [19, 25, 16],  # Dark green trees
+            [26, 33, 23],
+            [22, 30, 20],
+            [24, 35, 25]
         ])
         
+        # Expanded water colors
         self.water_colors = cp.array([
+            [40, 18, 4],   # Dark blue water
+            [39, 25, 6],
             [167, 225, 217],
             [67, 101, 97],
             [53, 83, 84],
             [47, 94, 100],
             [73, 131, 135]
         ])
+        
+        # Existing color arrays with optimized memory layout
+        self.shadow_colors = cp.asarray([
+            [31, 42, 76],
+            [58, 64, 92],
+            [15, 27, 56],
+            [21, 22, 50],
+            [76, 81, 99]
+        ], order='C')  # Use C-contiguous memory layout
+        
+        self.road_colors = cp.asarray([
+            [187, 182, 175],
+            [138, 138, 138], 
+            [142, 142, 129],
+            [202, 199, 189]
+        ], order='C')
 
-        # 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]
-        ])
+        # Output colors (BGR for OpenCV) - optimized memory layout
+        self.colors = {
+            'black': cp.asarray([0, 0, 0], order='C'),     # Shadows
+            'blue': cp.asarray([255, 0, 0], order='C'),    # Water
+            'dark_green': cp.asarray([0, 100, 0], order='C'),  # Trees
+            'light_green': cp.asarray([0, 255, 0], order='C'), # Grass
+            'gray': cp.asarray([128, 128, 128], order='C'),    # Roads
+            'brown': cp.asarray([0, 140, 255], order='C'),     # Terrain
+            'white': cp.asarray([255, 255, 255], order='C'),   # Buildings
+            'salmon': cp.asarray([128, 128, 255], order='C')   # Roofs
+        }
         
-        # 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 all color arrays to HSV space at initialization
+        self.initialize_hsv_colors()
+        
+        # Pre-compute kernels for morphological operations
+        self.cleanup_kernel = cp.ones((3, 3), dtype=bool)
+        self.cleanup_kernel[1, 1] = False
+        self.tree_kernel = cp.ones((5, 5), dtype=bool)
+        
+        # Optimization parameters
+        self.min_area = 1000
+        self.eps = 0.3
+        self.min_samples = 5
+    def initialize_hsv_colors(self):
+        """Initialize all HSV color spaces at once"""
+        color_arrays = {
+            'grass': self.grass_colors,
+            'tree': self.tree_colors,
+            'water': self.water_colors,
+            'shadow': self.shadow_colors,
+            'road': self.road_colors
         }
         
-        # Convert reference colors to HSV on GPU
-        self.shadow_colors_hsv = cp.asarray(cv2.cvtColor(
-            self.shadow_colors.get().reshape(-1, 1, 3).astype(np.uint8),
-            cv2.COLOR_RGB2HSV
-        ).reshape(-1, 3))
-        
-        self.road_colors_hsv = cp.asarray(cv2.cvtColor(
-            self.road_colors.get().reshape(-1, 1, 3).astype(np.uint8),
-            cv2.COLOR_RGB2HSV
-        ).reshape(-1, 3))
-        
-        self.water_colors_hsv = cp.asarray(cv2.cvtColor(
-            self.water_colors.get().reshape(-1, 1, 3).astype(np.uint8),
-            cv2.COLOR_RGB2HSV
-        ).reshape(-1, 3))
-        
-        # Normalize HSV values on GPU
-        for colors_hsv in [self.shadow_colors_hsv, self.road_colors_hsv, self.water_colors_hsv]:
-            colors_hsv[:, 0] = colors_hsv[:, 0] * 2
-            colors_hsv[:, 1:] = colors_hsv[:, 1:] / 255
-        
-        # Color tolerances
-        self.shadow_tolerance = {'hue': 15, 'sat': 0.15, 'val': 0.12}
-        self.road_tolerance = {'hue': 10, 'sat': 0.12, 'val': 0.15}
-        self.water_tolerance = {'hue': 20, 'sat': 0.15, 'val': 0.20}
-        
-        # 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
-            'salmon': cp.array([128, 128, 255])  # Roofs
+        self.hsv_colors = {}
+        self.tolerances = {
+            'grass': {'hue': 15, 'sat': 0.2, 'val': 0.15},
+            'tree': {'hue': 12, 'sat': 0.25, 'val': 0.15},
+            'water': {'hue': 25, 'sat': 0.2, 'val': 0.25},
+            'shadow': {'hue': 15, 'sat': 0.15, 'val': 0.12},
+            'road': {'hue': 10, 'sat': 0.12, 'val': 0.15}
         }
         
-        self.min_area_for_clustering = 1000
-        self.residential_height_factor = 0.6
-        self.isolation_threshold = 0.6
+        for name, colors in color_arrays.items():
+            hsv = cv2.cvtColor(colors.get().reshape(-1, 1, 3).astype(np.uint8),
+                             cv2.COLOR_RGB2HSV)
+            hsv_gpu = cp.asarray(hsv.reshape(-1, 3))
+            hsv_gpu[:, 0] = hsv_gpu[:, 0] * 2  # Scale hue to 0-360
+            hsv_gpu[:, 1:] = hsv_gpu[:, 1:] / 255  # Normalize S and V
+            self.hsv_colors[name] = hsv_gpu
 
     @staticmethod
+    @cp.fuse()  # Use CuPy's JIT compilation
     def gpu_color_distance_hsv(pixel_hsv, reference_hsv, tolerance):
-        """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]))
-        
-        # 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']
-            ),
-            val_diff <= tolerance['val']
-        )
+        """Optimized HSV color distance calculation using CuPy's JIT"""
+        h_diff = cp.minimum(cp.abs(pixel_hsv[0] - reference_hsv[0]),
+                          360 - cp.abs(pixel_hsv[0] - reference_hsv[0]))
+        s_diff = cp.abs(pixel_hsv[1] - reference_hsv[1])
+        v_diff = cp.abs(pixel_hsv[2] - reference_hsv[2])
+        
+        return (h_diff <= tolerance['hue']) & \
+               (s_diff <= tolerance['sat']) & \
+               (v_diff <= tolerance['val'])
+
+    def generate_tree_vertices(self, tree_mask, base_vertices):
+        """Generate randomized tree heights and positions"""
+        tree_positions = cp.where(tree_mask)
+        num_trees = len(tree_positions[0])
+        
+        # Random height variation for trees
+        tree_heights = cp.random.uniform(0.15, 0.25, num_trees)
+        
+        # Create vertex displacements for tree geometry
+        tree_vertices = base_vertices.copy()
+        tree_vertices[tree_positions] += cp.stack([
+            cp.zeros(num_trees),  # x offset
+            tree_heights,         # y offset (height)
+            cp.zeros(num_trees)   # z offset
+        ], axis=1)
+        
+        return tree_vertices
 
     def segment_image_gpu(self, img):
-        """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))
+        """Optimized GPU-accelerated image segmentation"""
+        # Transfer image to GPU with optimal memory layout
+        gpu_img = cp.asarray(img, order='C')
+        gpu_hsv = cp.asarray(cv2.cvtColor(img, cv2.COLOR_BGR2HSV), order='C')
         
         height, width = img.shape[:2]
-        output = cp.zeros_like(gpu_img)
-        
-        # Create a sliding window view for neighborhood analysis
-        pad = 2
-        gpu_hsv_pad = cp.pad(gpu_hsv, ((pad, pad), (pad, pad), (0, 0)), mode='edge')
+        output = cp.zeros_like(gpu_img, order='C')
         
-        # Prepare flattened HSV data
+        # Prepare HSV data
         hsv_pixels = gpu_hsv.reshape(-1, 3)
-        
-        # 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)
-        
-        # Color matching for predefined categories
-        for ref_hsv in self.shadow_colors_hsv:
-            temp_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:
-            temp_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)
-        
-        # 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
         v = v / 255
         
-        # Enhanced vegetation detection
-        vegetation_mask = ((h >= 40) & (h <= 150) & (s >= 0.15))
-        
-        # 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)
-        road_mask |= gray_mask & ~(shadow_mask | water_mask | vegetation_mask | terrain_mask | roof_mask)
-        
-        # Enhanced shadow detection
-        dark_mask = (v <= 0.3)
-        shadow_mask |= dark_mask & ~(water_mask | road_mask | roof_mask)
+        # Initialize masks with pre-allocated memory
+        masks = {
+            'shadow': cp.zeros(height * width, dtype=bool),
+            'road': cp.zeros(height * width, dtype=bool),
+            'water': cp.zeros(height * width, dtype=bool),
+            'grass': cp.zeros(height * width, dtype=bool),
+            'tree': cp.zeros(height * width, dtype=bool)
+        }
         
-        # Building mask (everything that's not another category)
-        building_mask = ~(shadow_mask | water_mask | road_mask | vegetation_mask | terrain_mask | roof_mask)
+        # Parallel color matching using CuPy's optimized operations
+        for category, hsv_refs in self.hsv_colors.items():
+            tolerance = self.tolerances[category]
+            for ref_hsv in hsv_refs:
+                masks[category] |= self.gpu_color_distance_hsv(
+                    cp.stack([h, s, v]),
+                    ref_hsv,
+                    tolerance
+                )
+        
+        # Optimized terrain and building detection
+        vegetation_mask = ((h >= 40) & (h <= 150) & (s >= 0.15))
+        terrain_mask = ((h >= 15) & (h <= 35) & (s >= 0.15) & (s <= 0.6))
+        building_mask = ~(masks['shadow'] | masks['water'] | masks['road'] | 
+                         masks['grass'] | masks['tree'] | vegetation_mask | 
+                         terrain_mask)
         
-        # Apply masks to create output
+        # Apply masks efficiently using CuPy's advanced indexing
         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']
+        for category, color_name in [
+            ('shadow', 'black'),
+            ('water', 'blue'),
+            ('grass', 'light_green'),
+            ('tree', 'dark_green'),
+            ('road', 'gray')
+        ]:
+            output_flat[masks[category]] = self.colors[color_name]
+        
         output_flat[terrain_mask] = self.colors['brown']
-        output_flat[roof_mask] = self.colors['salmon']
         output_flat[building_mask] = self.colors['white']
         
+        # Reshape and clean up
         segmented = output.reshape(height, width, 3)
+        segmented = self.apply_morphological_cleanup(segmented)
         
-        # Enhanced cleanup with roof consideration
-        kernel = cp.ones((3, 3), dtype=bool)
-        kernel[1, 1] = False
-        
-        # Two-pass cleanup
-        for _ in range(2):
+        return segmented
+
+    def apply_morphological_cleanup(self, segmented):
+        """Apply optimized morphological operations for cleanup"""
+        for _ in range(2):  # Two passes for better results
             for color_name, color_value in self.colors.items():
-                if cp.array_equal(color_value, self.colors['white']):
+                if color_name in ['white', 'dark_green']:  # Skip buildings and trees
                     continue
                 
                 color_mask = cp.all(segmented == color_value, axis=2)
-                dilated = binary_dilation(color_mask, structure=kernel)
+                dilated = binary_dilation(color_mask, structure=self.cleanup_kernel)
                 
                 building_pixels = cp.all(segmented == self.colors['white'], axis=2)
-                neighbor_count = binary_dilation(color_mask, structure=kernel).astype(int)
-                
-                # 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
+                neighbor_count = cp.sum(dilated)
                 
-                segmented[surrounded] = color_value
+                if neighbor_count > 5:
+                    segmented[building_pixels & dilated] = color_value
         
         return segmented
+
             
     def estimate_heights_gpu(self, img, segmented):
         """GPU-accelerated height estimation with roof consideration"""
@@ -294,14 +267,17 @@ class GPUSatelliteModelGenerator:
         return height_map.get() * 0.25
 
     def generate_mesh_gpu(self, height_map, texture_img):
-        """Generate 3D mesh using GPU-accelerated calculations"""
+        """Generate optimized 3D mesh with tree geometry"""
         height_map_gpu = cp.asarray(height_map)
         height, width = height_map.shape
         
-        # Generate vertex positions on GPU
+        # Generate base vertices
         x, z = cp.meshgrid(cp.arange(width), cp.arange(height))
         vertices = cp.stack([x, height_map_gpu * self.building_height, z], axis=-1)
-        vertices = vertices.reshape(-1, 3)
+        
+        # Detect tree areas and generate tree geometry
+        tree_mask = cp.all(texture_img == self.colors['dark_green'], axis=2)
+        vertices = self.generate_tree_vertices(tree_mask, vertices.reshape(-1, 3))
         
         # Normalize coordinates
         scale = max(width, height)
@@ -309,45 +285,52 @@ class GPUSatelliteModelGenerator:
         vertices[:, 2] = vertices[:, 2] / scale * 2 - (height / scale)
         vertices[:, 1] = vertices[:, 1] * 2 - 1
         
-        # Generate faces
+        # Generate optimized faces and UVs
+        faces = self.generate_faces_gpu(height, width)
+        uvs = self.generate_uvs_gpu(vertices, width, height)
+        
+        # Create textured mesh
+        return self.create_textured_mesh(vertices, faces, uvs, texture_img)
+
+    @staticmethod
+    def generate_faces_gpu(height, width):
+        """Generate optimized face indices"""
         i, j = cp.meshgrid(cp.arange(height-1), cp.arange(width-1), indexing='ij')
         v0 = (i * width + j).flatten()
         v1 = v0 + 1
         v2 = ((i + 1) * width + j).flatten()
         v3 = v2 + 1
         
-        faces = cp.vstack((
+        return cp.vstack((
             cp.column_stack((v0, v2, v1)),
             cp.column_stack((v1, v2, v3))
         ))
-        
-        # Generate UV coordinates
-        uvs = cp.zeros((vertices.shape[0], 2))
-        uvs[:, 0] = x.flatten() / (width - 1)
-        uvs[:, 1] = 1 - (z.flatten() / (height - 1))
-        
-        # Convert to CPU for mesh creation
-        vertices_cpu = vertices.get()
-        faces_cpu = faces.get()
-        uvs_cpu = uvs.get()
-        
-        # Create mesh
+
+    @staticmethod
+    def generate_uvs_gpu(vertices, width, height):
+        """Generate optimized UV coordinates"""
+        uvs = cp.zeros((vertices.shape[0], 2), order='C')
+        uvs[:, 0] = vertices[:, 0] / width
+        uvs[:, 1] = 1 - (vertices[:, 2] / height)
+        return uvs
+
+    @staticmethod
+    def create_textured_mesh(vertices, faces, uvs, texture_img):
+        """Create textured mesh with proper color conversion"""
         if len(texture_img.shape) == 3 and texture_img.shape[2] == 4:
             texture_img = cv2.cvtColor(texture_img, cv2.COLOR_BGRA2RGB)
         elif len(texture_img.shape) == 3:
             texture_img = cv2.cvtColor(texture_img, cv2.COLOR_BGR2RGB)
         
-        mesh = trimesh.Trimesh(
-            vertices=vertices_cpu,
-            faces=faces_cpu,
+        return trimesh.Trimesh(
+            vertices=vertices.get(),
+            faces=faces.get(),
             visual=trimesh.visual.TextureVisuals(
-                uv=uvs_cpu,
+                uv=uvs.get(),
                 image=Image.fromarray(texture_img)
             )
         )
         
-        return mesh
-        
 def generate_and_process_map(prompt: str) -> tuple[str | None, np.ndarray | None]:
     """Generate satellite image from prompt and convert to 3D model using GPU acceleration"""
     try: