video_depth.py

#!/usr/bin/env python
"""Turn any video into a FaceDepth depth-map video.

Decodes anything ffmpeg can read (mov, mp4, mkv, and the rest, up to 4K), runs every
frame through FaceDepth, and re-encodes. A progress bar reports frames per second and
ETA. Output is either the depth map alone or a side-by-side with the source on the left.

Frames stream through an ffmpeg pipe, so a long 4K clip never lands on disk as a frame
dump and memory stays flat.

Stability
---------
FaceDepth predicts each frame independently, which can make the colour mapping pulse
between frames. Two inference-time controls damp that:

  --range-ema     smooths the near/far normalisation range across frames. On by default.
                  Removes brightness pulsing with no ghosting. Leave it on.
  --smooth-depth  blends each depth frame with the previous one. Off by default.
                  Cuts residual per-pixel jitter, but ghosts behind fast motion.
                  0.3 is a reasonable starting point for handheld footage.

Setup
-----
    pip install torch torchvision opencv-python numpy
    git clone https://github.com/DepthAnything/Depth-Anything-V2 third_party/DepthAnythingV2

Download FaceDepth_step15792.pt from https://huggingface.co/a-ml/FaceDepth and pass it
with --ckpt. ffmpeg must be on PATH.

Examples
--------
    python video_depth.py --input clip.mov --output depth.mp4
    python video_depth.py --input clip.mkv --output sbs.mp4 --side-by-side --colormap turbo
    python video_depth.py --input 4k.mp4 --output out.mp4 --smooth-depth 0.3 --bf16

Runs on Apple silicon (mps), CUDA, or CPU, picked automatically.
"""
import argparse, json, subprocess, sys, time
from pathlib import Path

import cv2
import numpy as np
import torch

ROOT = Path(__file__).resolve().parent
for cand in (ROOT / "third_party" / "DepthAnythingV2", ROOT.parent / "third_party" / "DepthAnythingV2"):
    if cand.exists():
        sys.path.insert(0, str(cand))
        break
try:
    from depth_anything_v2.dpt import DepthAnythingV2
except ImportError:
    sys.exit("Could not import depth_anything_v2. Clone Depth-Anything-V2 into "
             "third_party/DepthAnythingV2 (see the setup notes at the top of this file).")

MEAN = torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1)
STD = torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1)
COLORMAPS = {"inferno": cv2.COLORMAP_INFERNO, "magma": cv2.COLORMAP_MAGMA,
             "turbo": cv2.COLORMAP_TURBO, "viridis": cv2.COLORMAP_VIRIDIS,
             "plasma": cv2.COLORMAP_PLASMA, "bone": cv2.COLORMAP_BONE, "gray": None}


def pick_device():
    if torch.backends.mps.is_available():
        return "mps"
    if torch.cuda.is_available():
        return "cuda"
    return "cpu"


def probe(path):
    out = subprocess.run(
        ["ffprobe", "-v", "error", "-select_streams", "v:0",
         "-show_entries", "stream=width,height,r_frame_rate,nb_frames,duration",
         "-of", "json", path], capture_output=True, text=True, check=True).stdout
    s = json.loads(out)["streams"][0]
    w, h = int(s["width"]), int(s["height"])
    num, den = s["r_frame_rate"].split("/")
    fps = float(num) / float(den)
    n = int(s["nb_frames"]) if s.get("nb_frames", "N/A") not in ("N/A", None) else \
        int(round(float(s.get("duration", 0)) * fps)) or 0
    return w, h, fps, n


def read_exact(pipe, n):
    buf = b""
    while len(buf) < n:
        chunk = pipe.read(n - len(buf))
        if not chunk:
            return None
        buf += chunk
    return buf


def fit14(w, h, res):
    """Scale so the longer side is about `res`, with both sides multiples of 14."""
    scale = res / max(w, h)
    return (max(14, int(round(w * scale / 14)) * 14),
            max(14, int(round(h * scale / 14)) * 14))


def bar(i, n, t0):
    frac = i / n if n else 0
    filled = int(40 * frac)
    el = time.time() - t0
    fps = i / max(el, 1e-6)
    eta = (n - i) / max(fps, 1e-6) if n else 0
    print(f"\r[{'#' * filled}{'-' * (40 - filled)}] {i}/{n} {100 * frac:5.1f}%  "
          f"{fps:4.1f} fps  ETA {eta:5.0f}s", end="", flush=True)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--input", required=True)
    ap.add_argument("--output", required=True)
    ap.add_argument("--ckpt", default="FaceDepth_step15792.pt")
    ap.add_argument("--res", type=int, default=910,
                    help="model longer-side resolution, rounded to a multiple of 14. "
                         "Higher is sharper and slower.")
    ap.add_argument("--side-by-side", action="store_true", help="source left, depth right")
    ap.add_argument("--colormap", choices=list(COLORMAPS), default="inferno")
    ap.add_argument("--invert", action="store_true", help="flip so near reads dark")
    ap.add_argument("--near-pct", type=float, default=2.0)
    ap.add_argument("--far-pct", type=float, default=98.0)
    ap.add_argument("--range-ema", type=float, default=0.85,
                    help="temporal smoothing of the near/far range, 0 disables")
    ap.add_argument("--smooth-depth", type=float, default=0.0,
                    help="temporal smoothing of depth itself, 0 disables, ghosts on motion")
    ap.add_argument("--bf16", action="store_true", help="roughly 2x faster, negligible quality cost")
    ap.add_argument("--keep-audio", action="store_true", default=True)
    ap.add_argument("--no-audio", dest="keep_audio", action="store_false")
    ap.add_argument("--crf", type=int, default=16, help="x264 quality, lower is better")
    args = ap.parse_args()

    inp = str(Path(args.input).expanduser())
    W, H, fps, N = probe(inp)
    mw, mh = fit14(W, H, args.res)
    outW, outH = (W * 2, H) if args.side_by_side else (W, H)
    print(f"input {W}x{H} @ {fps:.3f}fps, {N or '?'} frames  ->  model {mw}x{mh}  ->  "
          f"output {outW}x{outH} ({'side-by-side' if args.side_by_side else 'depth'})", flush=True)

    dev = pick_device()
    m = DepthAnythingV2(encoder="vitl", features=256, out_channels=[256, 512, 1024, 1024])
    ck = torch.load(args.ckpt, map_location="cpu", weights_only=True)
    m.load_state_dict(ck.get("ema_model") or ck.get("model") or ck)
    m = m.to(dev).eval()
    mean, std = MEAN.to(dev), STD.to(dev)
    print(f"loaded {Path(args.ckpt).name} on {dev}", flush=True)

    @torch.no_grad()
    def infer(rgb):
        x = torch.from_numpy(rgb).permute(2, 0, 1).unsqueeze(0).to(dev)
        if args.bf16 and dev != "cpu":
            with torch.autocast(dev, dtype=torch.bfloat16):
                d = m((x - mean) / std)
        else:
            d = m((x - mean) / std)
        return d.float().cpu().numpy()[0]

    dec = subprocess.Popen(["ffmpeg", "-v", "error", "-i", inp, "-f", "rawvideo",
                            "-pix_fmt", "bgr24", "-"], stdout=subprocess.PIPE, bufsize=10 ** 8)
    enc_cmd = ["ffmpeg", "-y", "-v", "error", "-f", "rawvideo", "-pix_fmt", "bgr24",
               "-s", f"{outW}x{outH}", "-r", f"{fps}", "-i", "-"]
    if args.keep_audio:
        enc_cmd += ["-i", inp, "-map", "0:v:0", "-map", "1:a:0?", "-c:a", "aac", "-shortest"]
    enc_cmd += ["-c:v", "libx264", "-crf", str(args.crf), "-pix_fmt", "yuv420p", args.output]
    enc = subprocess.Popen(enc_cmd, stdin=subprocess.PIPE, bufsize=10 ** 8)

    frame_bytes = W * H * 3
    ema_lo = ema_hi = ema_depth = None
    t0 = time.time()
    i = 0
    try:
        while True:
            raw = read_exact(dec.stdout, frame_bytes)
            if raw is None:
                break
            frame = np.frombuffer(raw, np.uint8).reshape(H, W, 3)
            small = cv2.resize(frame, (mw, mh), interpolation=cv2.INTER_AREA)
            rgb = cv2.cvtColor(small, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0
            disp = infer(rgb)  # inverse depth, larger is nearer

            if args.smooth_depth > 0:
                ema_depth = disp if ema_depth is None else \
                    (1 - args.smooth_depth) * disp + args.smooth_depth * ema_depth
                disp = ema_depth

            lo = np.percentile(disp, args.near_pct)
            hi = np.percentile(disp, args.far_pct)
            a = args.range_ema
            if a > 0:
                if ema_lo is None:
                    ema_lo, ema_hi = lo, hi
                else:
                    ema_lo = (1 - a) * lo + a * ema_lo
                    ema_hi = (1 - a) * hi + a * ema_hi
                lo, hi = ema_lo, ema_hi

            norm = np.clip((disp - lo) / (hi - lo + 1e-6), 0, 1)
            if args.invert:
                norm = 1 - norm
            gray = (norm * 255).astype(np.uint8)
            cmap = COLORMAPS[args.colormap]
            dvis = cv2.applyColorMap(gray, cmap) if cmap is not None else \
                cv2.cvtColor(gray, cv2.COLOR_GRAY2BGR)
            dvis = cv2.resize(dvis, (W, H), interpolation=cv2.INTER_CUBIC)

            out = np.hstack([frame, dvis]) if args.side_by_side else dvis
            enc.stdin.write(out.tobytes())

            i += 1
            if i % 5 == 0 or i == N:
                bar(i, N, t0)
    finally:
        print()
        if dec.stdout:
            dec.stdout.close()
        dec.wait()
        if enc.stdin:
            enc.stdin.close()
        enc.wait()
    print(f"done: {i} frames -> {args.output}  "
          f"({i / max(time.time() - t0, 1e-6):.1f} fps avg)", flush=True)


if __name__ == "__main__":
    main()