| 218 | |
| 219 | |
| 220 | def run_depth_midas_small(frames_rgb: List[np.ndarray], device: str) -> List[np.ndarray]: |
| 221 | try: |
| 222 | import torch |
| 223 | except ImportError as e: |
| 224 | raise ImportError("depth mode requires PyTorch (torch). Install the project requirements.txt.") from e |
| 225 | |
| 226 | if device.startswith("cuda") and torch.cuda.is_available(): |
| 227 | dev = torch.device(device) |
| 228 | elif device.startswith("cuda"): |
| 229 | logger.warning("CUDA requested for depth but not available; using CPU.") |
| 230 | dev = torch.device("cpu") |
| 231 | else: |
| 232 | dev = torch.device(device) |
| 233 | |
| 234 | logger.info("Loading MiDaS-small from torch.hub (first run may download weights)...") |
| 235 | midas = torch.hub.load("intel-isl/MiDaS", "MiDaS_small", trust_repo=True) |
| 236 | midas.to(dev).eval() |
| 237 | midas_transforms = torch.hub.load("intel-isl/MiDaS", "transforms") |
| 238 | transform = midas_transforms.small_transform |
| 239 | |
| 240 | out: List[np.ndarray] = [] |
| 241 | with torch.no_grad(): |
| 242 | for idx, frame in enumerate(frames_rgb): |
| 243 | h, w = frame.shape[:2] |
| 244 | im = Image.fromarray(frame) |
| 245 | inp = transform(im).to(dev) |
| 246 | pred = midas(inp) |
| 247 | if pred.ndim == 2: |
| 248 | pred = pred.unsqueeze(0) |
| 249 | pred = pred.unsqueeze(1) |
| 250 | pred = torch.nn.functional.interpolate( |
| 251 | pred, |
| 252 | size=(h, w), |
| 253 | mode="bicubic", |
| 254 | align_corners=False, |
| 255 | ) |
| 256 | d = pred.squeeze().cpu().numpy().astype(np.float32) |
| 257 | dmin, dmax = float(d.min()), float(d.max()) |
| 258 | if dmax - dmin < 1e-6: |
| 259 | u8 = np.zeros((h, w), dtype=np.uint8) |
| 260 | else: |
| 261 | u8 = ((d - dmin) / (dmax - dmin) * 255.0).clip(0, 255).astype(np.uint8) |
| 262 | rgb = np.stack([u8, u8, u8], axis=-1) |
| 263 | out.append(rgb) |
| 264 | if (idx + 1) % 20 == 0: |
| 265 | logger.info(f"Depth (MiDaS-small): {idx + 1}/{len(frames_rgb)} frames") |
| 266 | |
| 267 | del midas |
| 268 | if dev.type == "cuda": |
| 269 | try: |
| 270 | torch.cuda.empty_cache() |
| 271 | except Exception: |
| 272 | pass |
| 273 | return out |
| 274 | |
| 275 | |
| 276 | def _count_video_frames(path: str) -> int: |