Delete script.py

script.py

@@ -1,362 +0,0 @@
-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()