scripts/phase2_depth.py

import json
import os
import sys

import cv2
import numpy as np
import torch

ROOT = "/root/bdmc_pipeline"
FRAMES_DIR = f"{ROOT}/frames"
DEPTH_DIR = f"{ROOT}/depth"
CKPT = f"{ROOT}/depth_anything_v2/checkpoints/depth_anything_v2_metric_vkitti_vits.pth"
N_FRAMES = 300

sys.path.insert(0, f"{ROOT}/depth_anything_v2/metric_depth")
os.makedirs(DEPTH_DIR, exist_ok=True)

from depth_anything_v2.dpt import DepthAnythingV2

model_configs = {
    'vits': {'encoder': 'vits', 'features': 64, 'out_channels': [48, 96, 192, 384]},
}
encoder = 'vits'
max_depth = 80

print(f"Loading DepthAnything V2 Metric-Outdoor-Small (VKITTI, max_depth={max_depth})...")
model = DepthAnythingV2(**{**model_configs[encoder], 'max_depth': max_depth})
model.load_state_dict(torch.load(CKPT, map_location='cpu'))
device = "cuda" if torch.cuda.is_available() else "cpu"
model = model.to(device).eval()

depths = []
for i in range(1, N_FRAMES + 1):
    frame_path = f"{FRAMES_DIR}/f_{i:04d}.jpg"
    raw_img = cv2.imread(frame_path)
    if raw_img is None:
        raise FileNotFoundError(f"Missing frame: {frame_path}")

    with torch.no_grad():
        depth = model.infer_image(raw_img)

    np.save(f"{DEPTH_DIR}/f_{i:04d}.npy", depth.astype(np.float32))
    depths.append(depth)

    if i % 50 == 0 or i == 1:
        print(f"  Frame {i}/{N_FRAMES}: depth range [{depth.min():.2f}, {depth.max():.2f}] meters")

all_depths = np.stack(depths)
mean_d = float(all_depths.mean())
median_d = float(np.median(all_depths))
std_d = float(all_depths.std())

summary = {
    "method": "DepthAnything V2 Metric-Outdoor-Small (VKITTI pretrained, vits encoder)",
    "checkpoint": CKPT,
    "max_depth_m": max_depth,
    "n_frames": N_FRAMES,
    "depth_stats": {
        "mean_m": round(mean_d, 3),
        "median_m": round(median_d, 3),
        "std_m": round(std_d, 3),
    },
}
with open(f"{ROOT}/outputs/phase2_summary.json", "w") as f:
    json.dump(summary, f, indent=1)

print(f"\nPhase 2 complete. Depth maps saved to {DEPTH_DIR}/")
print(json.dumps(summary, indent=1))