03_depth_estimation.py

"""
03_depth_estimation.py — Monocular depth estimation on recording2.mp4

Model  : Depth Anything V2 Small (via HuggingFace transformers)
Output : outputs/depth_map.mp4  — side-by-side [original | depth colormap]
         Uses INFERNO colormap (warm = near, cool = far)

Notes:
- Fisheye distortion is present but Depth Anything V2 is robust to it
- All processing done on CPU (falls back gracefully from CUDA/MPS)
- Progress bar shows ETA
"""

import os
import sys
import numpy as np
import cv2
from tqdm import tqdm
from PIL import Image

# ── Paths ──────────────────────────────────────────────────────────────────────
BASE       = os.path.dirname(os.path.abspath(__file__))
VIDEO_PATH = os.path.join(BASE, 'data', 'recording2.mp4')
OUT_PATH   = os.path.join(BASE, 'outputs', 'depth_map.mp4')
os.makedirs(os.path.join(BASE, 'outputs'), exist_ok=True)

# ── Config ─────────────────────────────────────────────────────────────────────
MODEL_ID  = "depth-anything/Depth-Anything-V2-Small-hf"
OUT_W, OUT_H = 2560, 720   # side-by-side: 1280 cam + 1280 depth
PANEL_W   = 1280
PANEL_H   = 720


def load_depth_model():
    """Load Depth Anything V2 via transformers pipeline."""
    try:
        import torch
        from transformers import pipeline as hf_pipeline

        # Pick best available device
        if torch.cuda.is_available():
            device = 0       # CUDA GPU
            dtype  = torch.float16
        elif hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
            device = 'mps'
            dtype  = torch.float32
        else:
            device = -1      # CPU
            dtype  = torch.float32

        print(f"[03] Loading {MODEL_ID} on device={device} …")
        pipe = hf_pipeline(
            task="depth-estimation",
            model=MODEL_ID,
            device=device,
        )
        print("[03] Model loaded.")
        return pipe, 'transformers'

    except Exception as e:
        print(f"[03] transformers pipeline failed ({e}), falling back to MiDaS …")
        return load_midas_fallback()


def load_midas_fallback():
    """Fallback: MiDaS via torch.hub."""
    import torch
    print("[03] Loading MiDaS DPT_Large (fallback) …")
    model = torch.hub.load('intel-isl/MiDaS', 'DPT_Large')
    transforms = torch.hub.load('intel-isl/MiDaS', 'transforms')
    transform = transforms.dpt_transform
    model.eval()
    return (model, transform), 'midas'


def depth_to_colormap(depth_np: np.ndarray) -> np.ndarray:
    """Normalise depth map and apply INFERNO colormap → BGR uint8."""
    d_min, d_max = depth_np.min(), depth_np.max()
    if d_max > d_min:
        norm = ((depth_np - d_min) / (d_max - d_min) * 255).astype(np.uint8)
    else:
        norm = np.zeros_like(depth_np, dtype=np.uint8)
    return cv2.applyColorMap(norm, cv2.COLORMAP_INFERNO)


def draw_depth_labels(frame: np.ndarray, depth_np: np.ndarray) -> None:
    """Overlay NEAR/FAR labels and colourbar legend on the depth panel."""
    h, w = frame.shape[:2]
    # Colorbar strip on right edge
    bar_h, bar_w = h - 40, 18
    bar_x = w - bar_w - 10
    for row in range(bar_h):
        v = int((1 - row / bar_h) * 255)
        color = cv2.applyColorMap(np.array([[v]], dtype=np.uint8), cv2.COLORMAP_INFERNO)[0, 0]
        frame[20 + row, bar_x:bar_x + bar_w] = color
    cv2.putText(frame, 'NEAR', (bar_x - 36, 28),  cv2.FONT_HERSHEY_SIMPLEX, 0.45, (255,255,255), 1)
    cv2.putText(frame, 'FAR',  (bar_x - 30, h-22), cv2.FONT_HERSHEY_SIMPLEX, 0.45, (255,255,255), 1)
    # Min / max depth values
    cv2.putText(frame, f'min: {depth_np.min():.2f}', (10, h-30), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (200,200,200), 1)
    cv2.putText(frame, f'max: {depth_np.max():.2f}', (10, h-12), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (200,200,200), 1)
    # Title
    cv2.putText(frame, 'Depth Anything V2 (relative)', (10, 20),
                cv2.FONT_HERSHEY_SIMPLEX, 0.52, (79, 195, 247), 1, cv2.LINE_AA)


def process_transformers(pipe, cap, out, total_frames):
    """Run inference via HuggingFace transformers pipeline."""
    pbar = tqdm(total=total_frames, unit='frame', desc='Depth')
    frame_idx = 0
    while True:
        ret, bgr = cap.read()
        if not ret:
            break
        rgb_pil = Image.fromarray(cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB))

        result    = pipe(rgb_pil)
        depth_t   = result['predicted_depth']           # torch tensor
        depth_np  = depth_t.squeeze().numpy().astype(np.float32)

        # Resize panels
        cam_panel   = cv2.resize(bgr, (PANEL_W, PANEL_H))
        depth_color = depth_to_colormap(depth_np)
        depth_panel = cv2.resize(depth_color, (PANEL_W, PANEL_H))
        draw_depth_labels(depth_panel, depth_np)

        # Label separator line
        cv2.putText(cam_panel, 'Original', (10, 28), cv2.FONT_HERSHEY_SIMPLEX,
                    0.7, (255, 195, 79), 1, cv2.LINE_AA)
        # Frame counter
        cv2.putText(cam_panel, f'Frame {frame_idx:04d}/{total_frames}',
                    (10, PANEL_H - 12), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (200, 200, 200), 1)

        combined = np.hstack([cam_panel, depth_panel])
        out.write(combined)
        frame_idx += 1
        pbar.update(1)
    pbar.close()


def main():
    model_bundle, backend = load_depth_model()

    cap = cv2.VideoCapture(VIDEO_PATH)
    if not cap.isOpened():
        raise RuntimeError(f"Cannot open video: {VIDEO_PATH}")

    total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    fps   = cap.get(cv2.CAP_PROP_FPS)

    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    out    = cv2.VideoWriter(OUT_PATH, fourcc, fps, (OUT_W, OUT_H))
    print(f"[03] Processing {total} frames at {fps:.1f} fps → {OUT_PATH}")

    if backend == 'transformers':
        process_transformers(model_bundle, cap, out, total)
    else:
        # MiDaS fallback
        import torch
        model, transform = model_bundle
        pbar = tqdm(total=total, unit='frame', desc='Depth (MiDaS)')
        fi = 0
        while True:
            ret, bgr = cap.read()
            if not ret:
                break
            rgb = cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB)
            inp = transform(rgb).unsqueeze(0)
            with torch.no_grad():
                prediction = model(inp)
                depth_np = torch.nn.functional.interpolate(
                    prediction.unsqueeze(1), size=rgb.shape[:2],
                    mode='bicubic', align_corners=False
                ).squeeze().numpy()

            cam_panel   = cv2.resize(bgr, (PANEL_W, PANEL_H))
            depth_color = depth_to_colormap(depth_np)
            depth_panel = cv2.resize(depth_color, (PANEL_W, PANEL_H))
            draw_depth_labels(depth_panel, depth_np)
            combined = np.hstack([cam_panel, depth_panel])
            out.write(combined)
            fi += 1
            pbar.update(1)
        pbar.close()

    cap.release()
    out.release()
    print(f"[03] Done → {OUT_PATH}")


if __name__ == '__main__':
    main()