Spaces:
Running on Zero
Running on Zero
| """ | |
| M1 Perception Module - Monocular Depth Inference Engine | |
| """ | |
| import cv2 | |
| import logging | |
| import numpy as np | |
| from typing import Optional, Tuple | |
| from PIL import Image, ImageOps | |
| from transformers import pipeline | |
| class DepthEstimator: | |
| def __init__(self): | |
| self.pipe = None | |
| self.is_initialized = False | |
| def initialize(self, model_size: str = "Base") -> bool: | |
| """Loads the pretrained model weights into memory.""" | |
| # Map size strings to HuggingFace repo names | |
| self.model_size = model_size | |
| model_map = { | |
| "Small": "depth-anything/Depth-Anything-V2-Small-hf", | |
| "Base": "depth-anything/Depth-Anything-V2-Base-hf", | |
| "Large": "depth-anything/Depth-Anything-V2-Large-hf" | |
| } | |
| repo_id = model_map.get(model_size, model_map["Base"]) | |
| logging.info(f"[M1 MODEL] Initializing Depth Anything V2 {model_size} ({repo_id})...") | |
| try: | |
| self.pipe = pipeline(task="depth-estimation", model=repo_id) | |
| self.is_initialized = True | |
| return True | |
| except Exception as e: | |
| logging.error(f"[M1 MODEL] Failed to initialize model: {str(e)}") | |
| return False | |
| def predict_with_confidence(self, image: Image.Image, tta_mode: str = "2-view") -> Optional[Tuple[np.ndarray, np.ndarray]]: | |
| """ | |
| Executes depth inference with configurable Rotational Test-Time Augmentation (TTA). | |
| tta_mode: 'none' (1x), '2-view' (0° + 180°), or '4-view' (0°, 90°, 180°, 270°) | |
| Returns: (final_depth_array, confidence_map_array) | |
| """ | |
| if not self.is_initialized: | |
| logging.error("[M1 MODEL] Cannot predict: Model not initialized.") | |
| return None | |
| try: | |
| import time | |
| start_time = time.time() | |
| img_np = np.array(image) | |
| # Determine rotation steps based on mode | |
| if tta_mode == "4-view": | |
| rotations = [0, 1, 2, 3] | |
| elif tta_mode == "2-view": | |
| rotations = [0, 2] | |
| else: | |
| rotations = [0] | |
| preds = [] | |
| for k in rotations: | |
| # Rotate input image (k * 90 degrees) | |
| rotated_input = np.rot90(img_np, k) | |
| pil_rot = Image.fromarray(rotated_input) | |
| # Predict and extract the tensor | |
| res = self.pipe(pil_rot)["predicted_depth"] | |
| depth_arr = res.squeeze().cpu().numpy().astype(np.float32) | |
| # Rotate predicted depth back to canonical orientation | |
| canonical_depth = np.rot90(depth_arr, -k) | |
| preds.append(canonical_depth) | |
| # Consensus and variance mapping | |
| stacked = np.stack(preds, axis=0) | |
| final_depth = np.median(stacked, axis=0) | |
| # The standard deviation across rotations represents structural uncertainty | |
| if len(preds) > 1: | |
| uncertainty_map = np.std(stacked, axis=0) | |
| u_max = np.max(uncertainty_map) if np.max(uncertainty_map) > 0 else 1.0 | |
| confidence_map = 1.0 - (uncertainty_map / u_max) | |
| else: | |
| # Default to 1.0 confidence for single-pass inference | |
| confidence_map = np.ones_like(final_depth) | |
| # --------------------------------------------------------- | |
| # Calculate Terrain Bias B(D) using Morphological Opening | |
| # --------------------------------------------------------- | |
| # Calculate kernel size based on physical scale, not raw pixels | |
| gsd_meters_per_pixel = 1.0 # Default fallback, should be passed in from metadata | |
| max_building_width_meters = 60.0 # Assumes largest target building is 60m wide | |
| calculated_size = int(max_building_width_meters / gsd_meters_per_pixel) | |
| # OpenCV kernels must be odd numbers | |
| kernel_size = calculated_size if calculated_size % 2 != 0 else calculated_size + 1 | |
| kernel_size = max(3, kernel_size) # Ensure it doesn't collapse to 0 | |
| kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (kernel_size, kernel_size)) | |
| terrain_bias = cv2.morphologyEx(final_depth, cv2.MORPH_OPEN, kernel) | |
| height_map = np.maximum(final_depth - terrain_bias, 0.0) | |
| # --------------------------------------------------------- | |
| # Calculate latency in milliseconds | |
| inference_ms = round((time.time() - start_time) * 1000, 2) | |
| # Determine execution device from HuggingFace pipeline | |
| device_used = str(self.pipe.device) | |
| return { | |
| "relative_depth": final_depth, | |
| "relative_height": height_map, | |
| "quality_map": confidence_map, | |
| "metadata": { | |
| "model": f"Depth-Anything-V2-{self.model_size}", | |
| "tta_mode": tta_mode, | |
| "device": device_used, | |
| "inference_ms": inference_ms, | |
| "input_size": (img_np.shape[0], img_np.shape[1]), | |
| "dtype": str(final_depth.dtype) | |
| } | |
| } | |
| except Exception as e: | |
| logging.error(f"[M1 MODEL] Inference failed: {str(e)}") | |
| return None | |