Spaces:
Running on Zero
Running on Zero
| """ | |
| Module: depth.sat2map_refiner | |
| Description: sat2Map-inspired Architectural Roof Regularization, Planar Segmentation, | |
| and 3D Point Cloud / Mesh Export Engine (ACM CPS 2024 sat2Map methodology). | |
| Key Capabilities: | |
| 1. Structure-Preserving Roof Regularization: | |
| - Preserves sharp building eaves, ridgelines, and wall boundaries using edge-aware bilateral guidance. | |
| - Planar facet flattening: groups contiguous roof planes and regularizes normals to reduce spherical/blobby distortion. | |
| 2. Building Standardization & Height Parameter Recovery (sat2height inspired): | |
| - Computes spatial variance (sigma_x, sigma_y, sigma_z) and mean height (mu_h). | |
| 3. Planar Roof Extraction & Geometric Evaluation: | |
| - RANSAC-based planar segment detection. | |
| - Computes roof plane tilt angles, plane count, and planar quality index. | |
| 4. High-Fidelity 3D Model Exporters: | |
| - export_point_cloud_ply: Generates Stanford .ply colored point cloud with normals. | |
| - export_mesh_obj: Generates Wavefront .obj textured mesh with vertex normals & UV mapping. | |
| """ | |
| import os | |
| import struct | |
| import numpy as np | |
| from PIL import Image | |
| from typing import Optional, Dict, Any, Tuple, List | |
| from scipy.ndimage import gaussian_filter | |
| def regularize_roof_planes( | |
| height_map: np.ndarray, | |
| rgb: Optional[np.ndarray] = None, | |
| spatial_sigma: float = 3.0, | |
| range_sigma: float = 0.08, | |
| iterations: int = 2 | |
| ) -> np.ndarray: | |
| """ | |
| Edge-preserving roof planar regularization inspired by sat2pc residual refinement. | |
| Removes high-frequency monocular noise while preserving sharp architectural step edges. | |
| """ | |
| if height_map.ndim != 2: | |
| return height_map | |
| h_min = float(np.nanmin(height_map)) | |
| h_max = float(np.nanmax(height_map)) | |
| h_span = (h_max - h_min) if (h_max - h_min) > 1e-6 else 1.0 | |
| # Normalize to [0, 1] | |
| norm_h = np.clip((height_map - h_min) / h_span, 0.0, 1.0).astype(np.float32) | |
| refined = norm_h.copy() | |
| # Guided filtering if RGB guidance available, else bilateral approximation | |
| if rgb is not None and rgb.shape[:2] == height_map.shape: | |
| # Guidance luminance | |
| if rgb.ndim == 3: | |
| guidance = (0.299 * rgb[:, :, 0] + 0.587 * rgb[:, :, 1] + 0.114 * rgb[:, :, 2]) / 255.0 | |
| else: | |
| guidance = rgb / 255.0 | |
| guidance = guidance.astype(np.float32) | |
| else: | |
| guidance = norm_h.copy() | |
| # Iterative cross-bilateral / guided smoothing | |
| for _ in range(max(1, iterations)): | |
| mean_i = gaussian_filter(guidance, sigma=spatial_sigma) | |
| mean_p = gaussian_filter(refined, sigma=spatial_sigma) | |
| mean_ip = gaussian_filter(guidance * refined, sigma=spatial_sigma) | |
| cov_ip = mean_ip - mean_i * mean_p | |
| var_i = gaussian_filter(guidance * guidance, sigma=spatial_sigma) - mean_i * mean_i | |
| eps = (range_sigma ** 2) | |
| a = cov_ip / (var_i + eps) | |
| b = mean_p - a * mean_i | |
| mean_a = gaussian_filter(a, sigma=spatial_sigma) | |
| mean_b = gaussian_filter(b, sigma=spatial_sigma) | |
| refined = mean_a * guidance + mean_b | |
| # Gradient sharpening at dominant structural contours | |
| grad_y, grad_x = np.gradient(refined) | |
| grad_mag = np.sqrt(grad_x ** 2 + grad_y ** 2) | |
| edge_mask = grad_mag > np.percentile(grad_mag, 85) | |
| # Re-blend sharp transitions at structural edges | |
| refined = np.where(edge_mask, 0.65 * norm_h + 0.35 * refined, refined) | |
| # Restore original scale preserving min/max elevation datum | |
| restored = h_min + np.clip(refined, 0.0, 1.0) * h_span | |
| return restored.astype(np.float32) | |
| def extract_planar_segments( | |
| height_map: np.ndarray, | |
| gsd: float = 1.0, | |
| min_plane_size: int = 120, | |
| normal_thresh_deg: float = 18.0 | |
| ) -> Dict[str, Any]: | |
| """ | |
| Extracts planar roof facets from height raster using normal clustering & RANSAC plane fitting. | |
| Returns planar segmentation map, plane count, tilt angles, and quality stats. | |
| """ | |
| H, W = height_map.shape | |
| grad_y, grad_x = np.gradient(height_map.astype(np.float64), gsd, gsd) | |
| # Surface normals N = (-grad_x, -grad_y, 1.0) / len | |
| length = np.sqrt(grad_x ** 2 + grad_y ** 2 + 1.0) | |
| nx = -grad_x / length | |
| ny = -grad_y / length | |
| nz = 1.0 / length | |
| # Slope in degrees from vertical (z-axis) | |
| tilt_rad = np.arccos(np.clip(nz, -1.0, 1.0)) | |
| tilt_deg = np.degrees(tilt_rad) | |
| # Quantize normal angles into dominant roof facet bins (e.g. 8 azimuth sectors + flat) | |
| azimuth_rad = np.arctan2(ny, nx) | |
| azimuth_deg = (np.degrees(azimuth_rad) + 360.0) % 360.0 | |
| facet_map = np.zeros((H, W), dtype=np.int32) | |
| # Flat roofs / ground (< 6 deg tilt) | |
| flat_mask = tilt_deg < 6.0 | |
| facet_map[flat_mask] = 1 | |
| # Sloped roof facets categorized into 6 orientation bins | |
| sloped_mask = ~flat_mask | |
| az_bins = np.floor(azimuth_deg / 60.0).astype(np.int32) + 2 | |
| facet_map[sloped_mask] = az_bins[sloped_mask] | |
| # Filter tiny noisy regions | |
| unique_facets, counts = np.unique(facet_map, return_counts=True) | |
| valid_facets = [f for f, count in zip(unique_facets, counts) if count >= min_plane_size and f > 0] | |
| plane_tilts = [] | |
| for f_id in valid_facets: | |
| f_mask = (facet_map == f_id) | |
| avg_tilt = float(np.median(tilt_deg[f_mask])) | |
| plane_tilts.append(round(avg_tilt, 1)) | |
| return { | |
| "facet_map": facet_map, | |
| "num_planes": len(valid_facets), | |
| "plane_ids": [int(f) for f in valid_facets], | |
| "plane_tilts": plane_tilts, | |
| "mean_tilt_deg": round(float(np.mean(plane_tilts)), 1) if plane_tilts else 0.0, | |
| "nx": nx.astype(np.float32), | |
| "ny": ny.astype(np.float32), | |
| "nz": nz.astype(np.float32) | |
| } | |
| def export_point_cloud_ply( | |
| height_map: np.ndarray, | |
| rgb: Optional[np.ndarray], | |
| output_path: str, | |
| gsd_x: float = 1.0, | |
| gsd_y: float = 1.0, | |
| step: int = 1 | |
| ) -> str: | |
| """ | |
| Exports 3D Point Cloud in Stanford Binary PLY format with XYZ coordinates, | |
| RGB colors, and surface normals. | |
| """ | |
| H, W = height_map.shape | |
| os.makedirs(os.path.dirname(os.path.abspath(output_path)), exist_ok=True) | |
| # Subsample grid for point cloud | |
| rows = np.arange(0, H, step) | |
| cols = np.arange(0, W, step) | |
| grid_c, grid_r = np.meshgrid(cols, rows) | |
| xs = grid_c * gsd_x | |
| ys = (H - 1 - grid_r) * gsd_y # Invert Y for standard 3D cartesian coordinates | |
| zs = height_map[grid_r, grid_c] | |
| # Calculate normals | |
| grad_y, grad_x = np.gradient(height_map.astype(np.float64), gsd_y, gsd_x) | |
| length = np.sqrt(grad_x ** 2 + grad_y ** 2 + 1.0) | |
| nx = (-grad_x / length)[grid_r, grid_c] | |
| ny = (-grad_y / length)[grid_r, grid_c] | |
| nz = (1.0 / length)[grid_r, grid_c] | |
| if rgb is not None and rgb.shape[:2] == (H, W): | |
| r_chan = rgb[grid_r, grid_c, 0] | |
| g_chan = rgb[grid_r, grid_c, 1] | |
| b_chan = rgb[grid_r, grid_c, 2] | |
| else: | |
| # Default grayscale elevation ramp | |
| z_norm = np.clip((zs - np.nanmin(zs)) / (np.nanmax(zs) - np.nanmin(zs) + 1e-6), 0, 1) | |
| r_chan = (z_norm * 255).astype(np.uint8) | |
| g_chan = (z_norm * 255).astype(np.uint8) | |
| b_chan = (z_norm * 255).astype(np.uint8) | |
| num_vertices = len(grid_r.ravel()) | |
| with open(output_path, "wb") as f: | |
| header = ( | |
| f"ply\n" | |
| f"format binary_little_endian 1.0\n" | |
| f"element vertex {num_vertices}\n" | |
| f"property float x\n" | |
| f"property float y\n" | |
| f"property float z\n" | |
| f"property float nx\n" | |
| f"property float ny\n" | |
| f"property float nz\n" | |
| f"property uchar red\n" | |
| f"property uchar green\n" | |
| f"property uchar blue\n" | |
| f"end_header\n" | |
| ).encode("ascii") | |
| f.write(header) | |
| # Pack binary vertex data | |
| for x, y, z, n_x, n_y, n_z, r, g, b in zip( | |
| xs.ravel(), ys.ravel(), zs.ravel(), | |
| nx.ravel(), ny.ravel(), nz.ravel(), | |
| r_chan.ravel(), g_chan.ravel(), b_chan.ravel() | |
| ): | |
| f.write(struct.pack("<ffffffBBB", float(x), float(y), float(z), float(n_x), float(n_y), float(n_z), int(r), int(g), int(b))) | |
| return output_path | |
| def export_mesh_obj( | |
| height_map: np.ndarray, | |
| rgb: Optional[np.ndarray], | |
| output_obj_path: str, | |
| gsd_x: float = 1.0, | |
| gsd_y: float = 1.0, | |
| step: int = 2 | |
| ) -> Tuple[str, str, str]: | |
| """ | |
| Exports 3D Textured Mesh in Wavefront OBJ format with MTL material and JPEG texture map. | |
| """ | |
| H, W = height_map.shape | |
| base_dir = os.path.dirname(os.path.abspath(output_obj_path)) | |
| os.makedirs(base_dir, exist_ok=True) | |
| base_name = os.path.splitext(os.path.basename(output_obj_path))[0] | |
| mtl_filename = f"{base_name}.mtl" | |
| tex_filename = f"{base_name}_tex.jpg" | |
| mtl_path = os.path.join(base_dir, mtl_filename) | |
| tex_path = os.path.join(base_dir, tex_filename) | |
| # Save texture image | |
| if rgb is not None and rgb.shape[:2] == (H, W): | |
| pil_tex = Image.fromarray(rgb) | |
| else: | |
| z_norm = np.clip((height_map - np.nanmin(height_map)) / (np.nanmax(height_map) - np.nanmin(height_map) + 1e-6), 0, 1) | |
| pil_tex = Image.fromarray((z_norm * 255).astype(np.uint8)) | |
| pil_tex.save(tex_path, "JPEG", quality=90) | |
| # Write MTL file | |
| with open(mtl_path, "w", encoding="utf-8") as f_mtl: | |
| f_mtl.write( | |
| f"# DepthWizard Wavefront Material\n" | |
| f"newmtl DepthWizardMaterial\n" | |
| f"Ka 1.000 1.000 1.000\n" | |
| f"Kd 1.000 1.000 1.000\n" | |
| f"Ks 0.100 0.100 0.100\n" | |
| f"Ns 10.0\n" | |
| f"map_Kd {tex_filename}\n" | |
| ) | |
| # Build OBJ geometry grid | |
| rows = np.arange(0, H, step) | |
| cols = np.arange(0, W, step) | |
| n_rows = len(rows) | |
| n_cols = len(cols) | |
| with open(output_obj_path, "w", encoding="utf-8") as f_obj: | |
| f_obj.write(f"# DepthWizard 3D Reconstruction Model\n") | |
| f_obj.write(f"mtllib {mtl_filename}\n") | |
| f_obj.write(f"usemtl DepthWizardMaterial\n\n") | |
| # 1. Vertices (v x y z) | |
| for r_idx, r in enumerate(rows): | |
| y_val = (H - 1 - r) * gsd_y | |
| for c_idx, c in enumerate(cols): | |
| x_val = c * gsd_x | |
| z_val = float(height_map[r, c]) | |
| f_obj.write(f"v {x_val:.3f} {y_val:.3f} {z_val:.3f}\n") | |
| # 2. UV Texture coordinates (vt u v) | |
| for r_idx, r in enumerate(rows): | |
| v_uv = 1.0 - (r / (H - 1)) | |
| for c_idx, c in enumerate(cols): | |
| u_uv = c / (W - 1) | |
| f_obj.write(f"vt {u_uv:.4f} {v_uv:.4f}\n") | |
| # 3. Vertex Normals (vn nx ny nz) | |
| grad_y, grad_x = np.gradient(height_map.astype(np.float64), gsd_y, gsd_x) | |
| length = np.sqrt(grad_x ** 2 + grad_y ** 2 + 1.0) | |
| nx = (-grad_x / length) | |
| ny = (-grad_y / length) | |
| nz = (1.0 / length) | |
| for r_idx, r in enumerate(rows): | |
| for c_idx, c in enumerate(cols): | |
| f_obj.write(f"vn {nx[r, c]:.4f} {ny[r, c]:.4f} {nz[r, c]:.4f}\n") | |
| # 4. Faces (f v1/vt1/vn1 v2/vt2/vn2 v3/vt3/vn3) | |
| for r_idx in range(n_rows - 1): | |
| for c_idx in range(n_cols - 1): | |
| v_tl = r_idx * n_cols + c_idx + 1 | |
| v_tr = v_tl + 1 | |
| v_bl = (r_idx + 1) * n_cols + c_idx + 1 | |
| v_br = v_bl + 1 | |
| f_obj.write(f"f {v_tl}/{v_tl}/{v_tl} {v_bl}/{v_bl}/{v_bl} {v_tr}/{v_tr}/{v_tr}\n") | |
| f_obj.write(f"f {v_tr}/{v_tr}/{v_tr} {v_bl}/{v_bl}/{v_bl} {v_br}/{v_br}/{v_br}\n") | |
| return output_obj_path, mtl_path, tex_path | |