Update app.py

app.py

@@ -80,8 +80,161 @@ class GPUSatelliteModelGenerator:
         self.residential_height_factor = 0.6
         self.isolation_threshold = 0.6
 
-    # ... [Previous methods remain unchanged] ...
+    @staticmethod
+    def gpu_color_distance_hsv(pixel_hsv, reference_hsv, tolerance):
+        """GPU-accelerated HSV color distance calculation"""
+        pixel_h = pixel_hsv[0] * 2
+        pixel_s = pixel_hsv[1] / 255
+        pixel_v = pixel_hsv[2] / 255
+        
+        hue_diff = cp.minimum(cp.abs(pixel_h - reference_hsv[0]),
+                            360 - cp.abs(pixel_h - reference_hsv[0]))
+        sat_diff = cp.abs(pixel_s - reference_hsv[1])
+        val_diff = cp.abs(pixel_v - reference_hsv[2])
+        
+        return cp.logical_and(
+            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"""
+        # Transfer image to GPU
+        gpu_img = cp.asarray(img)
+        gpu_hsv = cp.asarray(cv2.cvtColor(img, cv2.COLOR_BGR2HSV))
+        
+        height, width = img.shape[:2]
+        output = cp.zeros_like(gpu_img)
+        
+        # Vectorized color matching on GPU
+        hsv_pixels = gpu_hsv.reshape(-1, 3)
+        
+        # Create masks for each category
+        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 ref_hsv in self.shadow_colors_hsv:
+            shadow_mask |= self.gpu_color_distance_hsv(hsv_pixels.T, ref_hsv, self.shadow_tolerance)
+            
+        for ref_hsv in self.road_colors_hsv:
+            road_mask |= self.gpu_color_distance_hsv(hsv_pixels.T, ref_hsv, self.road_tolerance)
+            
+        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
+        h, s, v = hsv_pixels.T
+        h = h * 2  # Convert to 0-360 range
+        s = s / 255
+        v = v / 255
+        
+        vegetation_mask = (h >= 40) & (h <= 150) & (s >= 0.15)
+        building_mask = ~(shadow_mask | water_mask | road_mask | vegetation_mask)
+        
+        output_flat[vegetation_mask] = self.colors['green']
+        output_flat[building_mask] = self.colors['white']
+        
+        return output.reshape(height, width, 3)
+        
+    def estimate_heights_gpu(self, img, segmented):
+        """GPU-accelerated height estimation"""
+        gpu_segmented = cp.asarray(segmented)
+        buildings_mask = cp.all(gpu_segmented == self.colors['white'], 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
+        max_area = cp.max(areas) if len(areas) > 0 else 1
+        
+        height_map = cp.zeros_like(labeled_array, dtype=cp.float32)
+        
+        # Process each building
+        for label in range(1, num_features + 1):
+            building_mask = (labeled_array == label)
+            if not cp.any(building_mask):
+                continue
+            
+            area = areas[label-1]
+            size_factor = 0.3 + 0.7 * (area / max_area)
+            
+            # 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
+        
+        return height_map.get() * 0.25
+
+    def generate_mesh_gpu(self, height_map, texture_img):
+        """Generate 3D mesh using GPU-accelerated calculations"""
+        height_map_gpu = cp.asarray(height_map)
+        height, width = height_map.shape
+        
+        # Generate vertex positions on GPU
+        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)
+        
+        # Normalize coordinates
+        scale = max(width, height)
+        vertices[:, 0] = vertices[:, 0] / scale * 2 - (width / scale)
+        vertices[:, 2] = vertices[:, 2] / scale * 2 - (height / scale)
+        vertices[:, 1] = vertices[:, 1] * 2 - 1
+        
+        # Generate faces
+        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((
+            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
+        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,
+            visual=trimesh.visual.TextureVisuals(
+                uv=uvs_cpu,
+                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: