Upload script.py

script.py

@@ -0,0 +1,362 @@
+import numpy as np
+import cv2
+import trimesh
+import argparse
+from PIL import Image
+from sklearn.cluster import KMeans
+
+class SatelliteModelGenerator:
+    def __init__(self, building_height=0.05):
+        self.building_height = building_height
+        
+        # Reference colors for segmentation (RGB)
+        self.shadow_colors = np.array([
+            [31, 42, 76],
+            [58, 64, 92],
+            [15, 27, 56],
+            [21, 22, 50],
+            [76, 81, 99]
+        ])
+        
+        self.road_colors = np.array([
+            [187, 182, 175],
+            [138, 138, 138], 
+            [142, 142, 129],
+            [202, 199, 189]
+        ])
+        
+        self.water_colors = np.array([
+            [167, 225, 217],
+            [67, 101, 97],
+            [53, 83, 84],
+            [47, 94, 100],
+            [73, 131, 135]
+        ])
+        
+        # Convert and normalize reference colors to HSV
+        self.shadow_colors_hsv = cv2.cvtColor(self.shadow_colors.reshape(-1, 1, 3).astype(np.uint8), 
+                                            cv2.COLOR_RGB2HSV).reshape(-1, 3).astype(float)
+        self.road_colors_hsv = cv2.cvtColor(self.road_colors.reshape(-1, 1, 3).astype(np.uint8), 
+                                          cv2.COLOR_RGB2HSV).reshape(-1, 3).astype(float)
+        self.water_colors_hsv = cv2.cvtColor(self.water_colors.reshape(-1, 1, 3).astype(np.uint8), 
+                                           cv2.COLOR_RGB2HSV).reshape(-1, 3).astype(float)
+        
+        # Normalize HSV values
+        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 from original segmenter
+        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}
+        
+        # Output colors (BGR for OpenCV)
+        self.colors = {
+            'black': np.array([0, 0, 0]),      # Shadows
+            'blue': np.array([255, 0, 0]),     # Water
+            'green': np.array([0, 255, 0]),    # Vegetation
+            'gray': np.array([128, 128, 128]), # Roads
+            'brown': np.array([0, 140, 255]),  # Terrain
+            'white': np.array([255, 255, 255]) # Buildings
+        }
+        
+        # Constants for height estimation
+        self.shadow_search_distance = 5
+        self.min_area_for_clustering = 1000
+        self.residential_height_factor = 0.6
+        self.isolation_threshold = 0.6
+
+    def color_distance_hsv(self, pixel_hsv, reference_hsv, tolerance):
+        """Calculate if a pixel is within tolerance of reference color in HSV space"""
+        pixel_h = float(pixel_hsv[0]) * 2
+        pixel_s = float(pixel_hsv[1]) / 255
+        pixel_v = float(pixel_hsv[2]) / 255
+        
+        hue_diff = min(abs(pixel_h - reference_hsv[0]),
+                      360 - abs(pixel_h - reference_hsv[0]))
+        sat_diff = abs(pixel_s - reference_hsv[1])
+        val_diff = abs(pixel_v - reference_hsv[2])
+        
+        return (hue_diff <= tolerance['hue'] and
+                sat_diff <= tolerance['sat'] and
+                val_diff <= tolerance['val'])
+
+    def get_dominant_surrounding_color(self, output, y, x):
+        """Determine dominant non-building color in neighborhood"""
+        height, width = output.shape[:2]
+        surroundings = []
+        
+        for dy in [-1, 0, 1]:
+            for dx in [-1, 0, 1]:
+                if dx == 0 and dy == 0:
+                    continue
+                
+                ny, nx = y + dy, x + dx
+                if 0 <= ny < height and 0 <= nx < width:
+                    pixel_color = tuple(output[ny, nx].tolist())
+                    if not np.array_equal(output[ny, nx], self.colors['white']):
+                        surroundings.append(pixel_color)
+        
+        if not surroundings:
+            return None
+        
+        surrounding_ratio = len(surroundings) / 8.0
+        
+        if surrounding_ratio >= self.isolation_threshold:
+            color_counts = {}
+            for color in surroundings:
+                color_str = str(color)
+                color_counts[color_str] = color_counts.get(color_str, 0) + 1
+            
+            most_common = max(color_counts.items(), key=lambda x: x[1])[0]
+            return np.array(eval(most_common))
+        
+        return None
+
+    def segment_image(self, img, window_size=5):
+        """Segment image using improved color detection"""
+        hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
+        output = np.zeros_like(img)
+        
+        pad = window_size // 2
+        hsv_pad = np.pad(hsv, ((pad, pad), (pad, pad), (0, 0)), mode='edge')
+        
+        height, width = img.shape[:2]
+        
+        # First pass: initial segmentation
+        for y in range(height):
+            for x in range(width):
+                window = hsv_pad[y:y+window_size, x:x+window_size]
+                center_hsv = window[pad, pad]
+                
+                is_shadow = any(self.color_distance_hsv(center_hsv, ref_hsv, self.shadow_tolerance)
+                              for ref_hsv in self.shadow_colors_hsv)
+                
+                is_road = any(self.color_distance_hsv(center_hsv, ref_hsv, self.road_tolerance)
+                            for ref_hsv in self.road_colors_hsv)
+                
+                is_water = any(self.color_distance_hsv(center_hsv, ref_hsv, self.water_tolerance)
+                             for ref_hsv in self.water_colors_hsv)
+                
+                if is_shadow:
+                    output[y, x] = self.colors['black']
+                elif is_water:
+                    output[y, x] = self.colors['blue']
+                elif is_road:
+                    output[y, x] = self.colors['gray']
+                else:
+                    h, s, v = center_hsv
+                    h = float(h) * 2  # Convert to 0-360 range
+                    s = float(s) / 255
+                    v = float(v) / 255
+                    
+                    # Check for pinkish building tones (around red hue with specific saturation)
+                    is_pinkish = (
+                        ((h >= 340 or h <= 15) and  # Red-pink hue range
+                         0.2 <= s <= 0.6 and        # Moderate saturation
+                         0.3 <= v <= 0.7)           # Moderate brightness
+                    )
+                    
+                    # Vegetation detection (green)
+                    is_vegetation = (
+                        40 <= h <= 150 and 
+                        s >= 0.15
+                    )
+                    
+                    # Soil/dirt detection (yellow-brown, avoiding pinkish tones)
+                    is_soil = (
+                        15 <= h <= 45 and           # Yellow-brown hue range
+                        0.15 <= s <= 0.45 and       # Lower saturation for dirt
+                        not is_pinkish              # Exclude pinkish tones
+                    )
+                    
+                    if is_pinkish:
+                        output[y, x] = self.colors['white']  # Buildings
+                    elif is_vegetation:
+                        output[y, x] = self.colors['green']  # Vegetation
+                    elif is_soil:
+                        output[y, x] = self.colors['brown']  # Soil/dirt
+                    else:
+                        # Default to building for light-colored surfaces
+                        output[y, x] = self.colors['white']
+        
+        # Second pass: handle isolated building pixels
+        final_output = output.copy()
+        for y in range(height):
+            for x in range(width):
+                if np.array_equal(output[y, x], self.colors['white']):
+                    dominant_color = self.get_dominant_surrounding_color(output, y, x)
+                    if dominant_color is not None:
+                        final_output[y, x] = dominant_color
+        
+        return final_output
+
+    def estimate_heights(self, img, segmented):
+        """Estimate building heights"""
+        buildings_mask = np.all(segmented == self.colors['white'], axis=2)
+        shadows_mask = np.all(segmented == self.colors['black'], axis=2)
+        
+        num_buildings, labels = cv2.connectedComponents(buildings_mask.astype(np.uint8))
+        
+        areas = np.bincount(labels.flatten())[1:]  # Skip background
+        max_area = np.max(areas) if len(areas) > 0 else 1
+        
+        height_map = np.zeros_like(labels, dtype=np.float32)
+        
+        for label in range(1, num_buildings):
+            building_mask = (labels == label)
+            if not np.any(building_mask):
+                continue
+            
+            area = areas[label-1]
+            size_factor = 0.3 + 0.7 * (area / max_area)
+            
+            dilated = cv2.dilate(building_mask.astype(np.uint8), np.ones((5,5), np.uint8))
+            shadow_ratio = np.sum(dilated & shadows_mask) / np.sum(dilated)
+            shadow_factor = 0.2 + 0.8 * shadow_ratio
+            
+            if area >= self.min_area_for_clustering:
+                building_intensities = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)[building_mask]
+                kmeans = KMeans(n_clusters=2, random_state=42)
+                clusters = kmeans.fit_predict(building_intensities.reshape(-1, 1))
+                cluster_means = [building_intensities[clusters == i].mean() for i in range(2)]
+                height_factor = self.residential_height_factor if cluster_means[0] > cluster_means[1] else 1.0
+            else:
+                height_factor = 1.0
+            
+            final_height = size_factor * shadow_factor * height_factor
+            height_map[building_mask] = final_height
+        
+        return height_map * 0.15
+
+    def generate_mesh(self, height_map, texture_img, add_walls=True):
+        """Generate 3D mesh"""
+        height, width = height_map.shape
+        
+        x, z = np.meshgrid(np.arange(width), np.arange(height))
+        vertices = np.stack([x, height_map * self.building_height, z], axis=-1)
+        vertices = vertices.reshape(-1, 3)
+        
+        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
+        
+        i, j = np.meshgrid(np.arange(height-1), np.arange(width-1), indexing='ij')
+        v0 = (i * width + j).flatten()
+        v1 = v0 + 1
+        v2 = ((i + 1) * width + j).flatten()
+        v3 = v2 + 1
+        
+        faces = np.vstack((
+            np.column_stack((v0, v2, v1)),
+            np.column_stack((v1, v2, v3))
+        ))
+        
+        uvs = np.zeros((vertices.shape[0], 2))
+        uvs[:, 0] = x.flatten() / (width - 1)
+        uvs[:, 1] = 1 - (z.flatten() / (height - 1))
+        
+        if len(texture_img.shape) == 3:
+            if texture_img.shape[2] == 4:
+                texture_img = cv2.cvtColor(texture_img, cv2.COLOR_BGRA2RGB)
+            else:
+                texture_img = cv2.cvtColor(texture_img, cv2.COLOR_BGR2RGB)
+        
+        mesh = trimesh.Trimesh(
+            vertices=vertices,
+            faces=faces,
+            visual=trimesh.visual.TextureVisuals(
+                uv=uvs,
+                image=Image.fromarray(texture_img)
+            )
+        )
+        
+        if add_walls:
+            mesh = self._add_walls(mesh, height_map)
+        
+        return mesh
+
+    def _add_walls(self, mesh, height_map):
+        """Add vertical walls at building edges"""
+        edges = cv2.Canny(height_map.astype(np.uint8) * 255, 100, 200)
+        height, width = height_map.shape
+        scale = max(width, height)
+        
+        edge_coords = np.column_stack(np.where(edges > 0))
+        if len(edge_coords) == 0:
+            return mesh
+        
+        valid_mask = (edge_coords[:, 0] < height - 1) & (edge_coords[:, 1] < width - 1)
+        edge_coords = edge_coords[valid_mask]
+        
+        if len(edge_coords) == 0:
+            return mesh
+        
+        y, x = edge_coords.T
+        heights = height_map[y, x]
+        
+        top_front = np.column_stack([x, heights * self.building_height, y])
+        top_back = np.column_stack([x + 1, heights * self.building_height, y])
+        bottom_front = np.column_stack([x, np.zeros_like(heights), y])
+        bottom_back = np.column_stack([x + 1, np.zeros_like(heights), y])
+        
+        for vertices in [top_front, top_back, bottom_front, bottom_back]:
+            vertices[:, 0] = vertices[:, 0] / scale * 2 - (width / scale)
+            vertices[:, 2] = vertices[:, 2] / scale * 2 - (height / scale)
+            vertices[:, 1] = vertices[:, 1] * 2 - 1
+        
+        new_vertices = np.vstack([top_front, top_back, bottom_front, bottom_back])
+        vertex_count = len(edge_coords)
+        
+        indices = np.arange(4 * vertex_count).reshape(-1, 4)
+        new_faces = np.vstack([
+            np.column_stack([indices[:, 0], indices[:, 2], indices[:, 1]]),
+            np.column_stack([indices[:, 1], indices[:, 2], indices[:, 3]])
+        ])
+        
+        base_vertex_count = len(mesh.vertices)
+        mesh.vertices = np.vstack((mesh.vertices, new_vertices))
+        mesh.faces = np.vstack((mesh.faces, new_faces + base_vertex_count))
+        
+        return mesh
+
+def main():
+    parser = argparse.ArgumentParser(description='Generate 3D mesh from satellite image')
+    parser.add_argument('input_image', help='Path to satellite image')
+    parser.add_argument('output_mesh', help='Path for output GLB file')
+    parser.add_argument('--segmented_output', help='Optional path to save segmented image')
+    parser.add_argument('--height', type=float, default=0.09, help='Height of buildings (default: 0.09)')
+    parser.add_argument('--no_walls', action='store_true', help='Skip generating vertical walls')
+    parser.add_argument('--window_size', type=int, default=5, help='Window size for segmentation analysis')
+    
+    args = parser.parse_args()
+    
+    # Load image
+    img = cv2.imread(args.input_image)
+    if img is None:
+        raise ValueError(f"Could not read image at {args.input_image}")
+    
+    generator = SatelliteModelGenerator(building_height=args.height)
+    
+    # Process image
+    print("Segmenting image...")
+    segmented_img = generator.segment_image(img, args.window_size)
+    
+    print("Estimating heights...")
+    height_map = generator.estimate_heights(img, segmented_img)
+    
+    # Save segmented image if requested
+    if args.segmented_output:
+        cv2.imwrite(args.segmented_output, segmented_img)
+        print(f"Segmented image saved to {args.segmented_output}")
+    
+    # Generate and save mesh
+    print("Generating mesh...")
+    mesh = generator.generate_mesh(height_map, img, add_walls=not args.no_walls)
+    mesh.export(args.output_mesh)
+    print(f"Mesh exported to {args.output_mesh}")
+
+if __name__ == "__main__":
+    main()