Spaces:
Running on Zero
Running on Zero
| """ | |
| Module: depth.tiled_inference | |
| Description: High-resolution overlapping tiled monocular depth inference using Depth Anything V2. | |
| """ | |
| import os | |
| import time | |
| import torch | |
| import numpy as np | |
| from PIL import Image | |
| from typing import Union, Optional, Tuple, Dict, Any | |
| from depth.depth_model import load_model, get_device | |
| def _create_2d_window(height: int, width: int, window_type: str = "hann") -> np.ndarray: | |
| if window_type == "hann": | |
| wy = np.sin(np.linspace(np.pi / (2 * height), np.pi - np.pi / (2 * height), height)) ** 2 | |
| wx = np.sin(np.linspace(np.pi / (2 * width), np.pi - np.pi / (2 * width), width)) ** 2 | |
| w2d = np.outer(wy, wx) | |
| elif window_type == "pyramid": | |
| wy = 1.0 - np.abs(np.linspace(-1, 1, height)) | |
| wx = 1.0 - np.abs(np.linspace(-1, 1, width)) | |
| w2d = np.outer(wy, wx) | |
| else: | |
| w2d = np.ones((height, width), dtype=np.float32) | |
| w2d = np.clip(w2d, 1e-4, 1.0).astype(np.float32) | |
| return w2d | |
| def estimate_depth_tiled( | |
| image: Union[str, np.ndarray, Image.Image], | |
| tile_size: int = 512, | |
| overlap: float = 0.25, | |
| batch_size: int = 8, | |
| window_type: str = "hann", | |
| device: Optional[Union[str, torch.device]] = None | |
| ) -> Tuple[np.ndarray, Dict[str, Any]]: | |
| start_time = time.time() | |
| if isinstance(image, str): | |
| if not os.path.exists(image): | |
| raise FileNotFoundError(f"Image file not found: {image}") | |
| pil_img = Image.open(image).convert("RGB") | |
| rgb_arr = np.array(pil_img, dtype=np.uint8) | |
| elif isinstance(image, np.ndarray): | |
| arr = np.squeeze(image) | |
| if arr.dtype != np.uint8: | |
| if arr.max() <= 1.0: | |
| arr = (arr * 255.0).astype(np.uint8) | |
| else: | |
| arr = np.clip(arr, 0, 255).astype(np.uint8) | |
| if arr.ndim == 2: | |
| arr = np.stack([arr] * 3, axis=-1) | |
| rgb_arr = arr | |
| elif isinstance(image, Image.Image): | |
| rgb_arr = np.array(image.convert("RGB"), dtype=np.uint8) | |
| else: | |
| raise TypeError(f"Unsupported input image type: {type(image)}") | |
| H, W = rgb_arr.shape[:2] | |
| stride = max(1, int(round(tile_size * (1.0 - overlap)))) | |
| y_starts = list(range(0, max(1, H - tile_size + 1), stride)) | |
| if len(y_starts) == 0 or (y_starts[-1] + tile_size < H): | |
| y_starts.append(max(0, H - tile_size)) | |
| x_starts = list(range(0, max(1, W - tile_size + 1), stride)) | |
| if len(x_starts) == 0 or (x_starts[-1] + tile_size < W): | |
| x_starts.append(max(0, W - tile_size)) | |
| y_starts = sorted(list(set(y_starts))) | |
| x_starts = sorted(list(set(x_starts))) | |
| tiles = [] | |
| for y0 in y_starts: | |
| y1 = min(H, y0 + tile_size) | |
| y0_eff = max(0, y1 - tile_size) | |
| for x0 in x_starts: | |
| x1 = min(W, x0 + tile_size) | |
| x0_eff = max(0, x1 - tile_size) | |
| tiles.append((y0_eff, y1, x0_eff, x1)) | |
| num_tiles = len(tiles) | |
| model, processor, target_device = load_model(device=device) | |
| base_window = _create_2d_window(tile_size, tile_size, window_type=window_type) | |
| accum_depth = np.zeros((H, W), dtype=np.float64) | |
| accum_weight = np.zeros((H, W), dtype=np.float64) | |
| if torch.cuda.is_available() and target_device.type == "cuda": | |
| torch.cuda.reset_peak_memory_stats(target_device) | |
| for b_idx in range(0, num_tiles, batch_size): | |
| batch_coords = tiles[b_idx : b_idx + batch_size] | |
| batch_pil = [ | |
| Image.fromarray(rgb_arr[y0:y1, x0:x1]) | |
| for (y0, y1, x0, x1) in batch_coords | |
| ] | |
| inputs = processor(images=batch_pil, return_tensors="pt").to(target_device) | |
| with torch.no_grad(): | |
| outputs = model(**inputs) | |
| post_processed = processor.post_process_depth_estimation( | |
| outputs, | |
| target_sizes=[(y1 - y0, x1 - x0) for (y0, y1, x0, x1) in batch_coords] | |
| ) | |
| for (y0, y1, x0, x1), pred in zip(batch_coords, post_processed): | |
| t_depth = pred["predicted_depth"].cpu().numpy().astype(np.float64) | |
| th, tw = y1 - y0, x1 - x0 | |
| if (th, tw) == (tile_size, tile_size): | |
| w = base_window | |
| else: | |
| w = _create_2d_window(th, tw, window_type=window_type) | |
| accum_depth[y0:y1, x0:x1] += t_depth * w | |
| accum_weight[y0:y1, x0:x1] += w | |
| valid_weight = accum_weight > 0 | |
| final_depth = np.zeros((H, W), dtype=np.float32) | |
| final_depth[valid_weight] = (accum_depth[valid_weight] / accum_weight[valid_weight]).astype(np.float32) | |
| elapsed = time.time() - start_time | |
| peak_gpu_mb = 0.0 | |
| if torch.cuda.is_available() and target_device.type == "cuda": | |
| peak_gpu_mb = float(torch.cuda.max_memory_allocated(target_device) / (1024 * 1024)) | |
| info = { | |
| "image_shape": [H, W, 3], | |
| "tile_size": tile_size, | |
| "overlap": overlap, | |
| "stride": stride, | |
| "grid_y": len(y_starts), | |
| "grid_x": len(x_starts), | |
| "total_tiles": num_tiles, | |
| "batch_size": batch_size, | |
| "window_type": window_type, | |
| "elapsed_seconds": elapsed, | |
| "peak_gpu_memory_mb": peak_gpu_mb, | |
| "device": str(target_device) | |
| } | |
| return final_depth, info | |
| # Alias for backwards compatibility | |
| run_tiled_inference = estimate_depth_tiled | |