script.py

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()