Spaces:
Running on Zero
Running on Zero
File size: 4,929 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 129 130 | """
Module: depth.depth_model
Description: Monocular Relative Depth Estimation using Depth Anything V2.
Depth Convention:
Depth Anything V2 outputs scale- and shift-invariant relative inverse depth (disparity).
- Higher numerical values correspond to regions closer to the optical sensor (i.e. elevated building rooftops).
- Lower numerical values correspond to regions farther from the optical sensor (i.e. ground-level terrain).
- Output is a 2D float32 NumPy array preserving the exact spatial dimensions (H, W) of the input RGB image.
"""
import os
import torch
import numpy as np
from PIL import Image
from typing import Union, Tuple, Optional
from transformers import AutoImageProcessor, AutoModelForDepthEstimation
DEFAULT_HF_MODEL_ID = "depth-anything/Depth-Anything-V2-Small-hf"
ENV_CHECKPOINT_VAR = "DEPTH_MODEL_CHECKPOINT"
_CACHED_MODEL = None
_CACHED_PROCESSOR = None
_CACHED_DEVICE = None
_CACHED_PATH = None
def get_device() -> torch.device:
if torch.cuda.is_available():
return torch.device("cuda")
return torch.device("cpu")
def load_model(
checkpoint_path: Optional[str] = None,
device: Optional[Union[str, torch.device]] = None
) -> Tuple[AutoModelForDepthEstimation, AutoImageProcessor, torch.device]:
global _CACHED_MODEL, _CACHED_PROCESSOR, _CACHED_DEVICE, _CACHED_PATH
if device is None:
target_device = get_device()
else:
target_device = torch.device(device)
if checkpoint_path is None:
checkpoint_path = os.environ.get(ENV_CHECKPOINT_VAR, DEFAULT_HF_MODEL_ID)
if _CACHED_MODEL is not None and _CACHED_DEVICE == target_device and _CACHED_PATH == checkpoint_path:
return _CACHED_MODEL, _CACHED_PROCESSOR, _CACHED_DEVICE
if checkpoint_path != DEFAULT_HF_MODEL_ID and not os.path.exists(checkpoint_path):
raise FileNotFoundError(
f"Depth Anything V2 checkpoint not found at: '{checkpoint_path}'.\n"
f"Please place the required weights in '{checkpoint_path}' or set the "
f"environment variable '{ENV_CHECKPOINT_VAR}' to a valid local path or HF model ID."
)
print(f"[DepthModel] Loading Depth Anything V2 from: '{checkpoint_path}' on device: {target_device}...")
try:
processor = AutoImageProcessor.from_pretrained(checkpoint_path)
model = AutoModelForDepthEstimation.from_pretrained(checkpoint_path).to(target_device)
model.eval()
except Exception as e:
raise RuntimeError(
f"Failed to load Depth Anything V2 from '{checkpoint_path}': {e}\n"
f"Ensure dependencies are installed and the checkpoint path is valid."
) from e
_CACHED_MODEL = model
_CACHED_PROCESSOR = processor
_CACHED_DEVICE = target_device
_CACHED_PATH = checkpoint_path
return model, processor, target_device
def estimate_depth(
image_input: Union[str, np.ndarray, Image.Image],
model: Optional[AutoModelForDepthEstimation] = None,
processor: Optional[AutoImageProcessor] = None,
device: Optional[Union[str, torch.device]] = None
) -> np.ndarray:
if isinstance(image_input, str):
if not os.path.exists(image_input):
raise FileNotFoundError(f"Image file not found: '{image_input}'")
pil_img = Image.open(image_input).convert("RGB")
elif isinstance(image_input, np.ndarray):
arr = np.squeeze(image_input)
if arr.dtype != np.uint8:
if arr.max() <= 1.0:
arr = (arr * 255.0).astype(np.uint8)
elif arr.max() > 255.0:
from data_io.cartosat_adapter import stretch_satellite_dynamic_range
arr = stretch_satellite_dynamic_range(arr)
else:
arr = np.clip(arr, 0, 255).astype(np.uint8)
if arr.ndim == 2:
arr = np.stack([arr]*3, axis=-1)
pil_img = Image.fromarray(arr).convert("RGB")
elif isinstance(image_input, Image.Image):
pil_img = image_input.convert("RGB")
else:
raise TypeError(f"Unsupported image input type: {type(image_input)}")
target_w, target_h = pil_img.size
if model is None or processor is None:
model, processor, device = load_model(device=device)
elif device is None:
device = next(model.parameters()).device
inputs = processor(images=pil_img, return_tensors="pt").to(device)
with torch.no_grad():
outputs = model(**inputs)
post_processed = processor.post_process_depth_estimation(
outputs,
target_sizes=[(target_h, target_w)]
)
depth_tensor = post_processed[0]["predicted_depth"]
depth_map = depth_tensor.cpu().numpy().astype(np.float32)
if depth_map.shape != (target_h, target_w):
raise ValueError(
f"Output depth shape {depth_map.shape} does not match input dimensions ({target_h}, {target_w})"
)
return depth_map
|