Spaces:
Paused
Paused
| """Depth Anything V2 inference with a lossless float32 prediction output.""" | |
| from __future__ import annotations | |
| import time | |
| from pathlib import Path | |
| import numpy as np | |
| import torch | |
| import torch.nn.functional as F | |
| from PIL import Image, ImageOps, UnidentifiedImageError | |
| from transformers import AutoImageProcessor, AutoModelForDepthEstimation | |
| MODEL_ID = "depth-anything/Depth-Anything-V2-Small-hf" | |
| def load_rgb(path: Path) -> Image.Image: | |
| if not path.is_file(): | |
| raise FileNotFoundError(f"Input image does not exist: {path}") | |
| try: | |
| with Image.open(path) as im: | |
| im = ImageOps.exif_transpose(im) | |
| im.load() | |
| return im.convert("RGB") | |
| except (UnidentifiedImageError, OSError, ValueError) as exc: | |
| raise ValueError(f"Unsupported or corrupted image {path}: {exc}") from exc | |
| def select_device() -> torch.device: | |
| if hasattr(torch.backends, "mps") and torch.backends.mps.is_available(): | |
| return torch.device("mps") | |
| return torch.device("cpu") | |
| def run_inference(input_path: Path, output_dir: Path, device: torch.device | None = None) -> tuple[Path, Path, float]: | |
| image = load_rgb(input_path) | |
| device = device or select_device() | |
| output_dir.mkdir(parents=True, exist_ok=True) | |
| stem = input_path.stem | |
| raw_path = output_dir / f"{stem}_depth_raw.npy" | |
| visual_path = output_dir / f"{stem}_depth.png" | |
| print(f"Device: {device}") | |
| print(f"Input: {input_path.resolve()}") | |
| print(f"Image dimensions: {image.width} x {image.height}") | |
| print(f"Model: {MODEL_ID}") | |
| started = time.perf_counter() | |
| try: | |
| processor = AutoImageProcessor.from_pretrained(MODEL_ID) | |
| model = AutoModelForDepthEstimation.from_pretrained(MODEL_ID).to(device).eval() | |
| except Exception as exc: | |
| raise RuntimeError(f"Could not load model {MODEL_ID}; check Hugging Face/network access: {exc}") from exc | |
| inputs = processor(images=image, return_tensors="pt") | |
| inputs = {key: value.to(device) for key, value in inputs.items()} | |
| with torch.inference_mode(): | |
| predicted = model(**inputs).predicted_depth | |
| predicted = F.interpolate(predicted.unsqueeze(1), size=(image.height, image.width), mode="bicubic", align_corners=False).squeeze().float() | |
| prediction = predicted.detach().cpu().numpy().astype(np.float32, copy=False) | |
| if prediction.shape != (image.height, image.width) or not np.isfinite(prediction).all(): | |
| raise RuntimeError(f"Model returned invalid prediction array: shape={prediction.shape}") | |
| elapsed = time.perf_counter() - started | |
| np.save(raw_path, prediction, allow_pickle=False) | |
| # Visualization only: robust percentile stretch; numeric evaluation uses the .npy. | |
| lo, hi = np.percentile(prediction, [2, 98]) | |
| shown = np.zeros_like(prediction, dtype=np.uint8) if hi <= lo else np.clip((prediction - lo) / (hi - lo) * 255, 0, 255).astype(np.uint8) | |
| Image.fromarray(shown, mode="L").save(visual_path) | |
| del model, processor, inputs, predicted | |
| if device.type == "mps": | |
| torch.mps.empty_cache() | |
| print("Inference complete") | |
| print(f"Raw prediction: {raw_path.resolve()}") | |
| print(f"Depth visualization: {visual_path.resolve()}") | |
| print(f"Inference time: {elapsed:.2f} seconds") | |
| return raw_path, visual_path, elapsed | |