Add video conversion script

video_depth.py

@@ -0,0 +1,227 @@
+#!/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()