Spaces:
Running on Zero
Running on Zero
File size: 5,319 Bytes
ad38b1a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 | """
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
|