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