#!/bin/sh
"exec" "`dirname $0`/../venvs/roto/bin/python" "$0" "$@"
"""vg-remove — erase objects/logos/watermarks from video (ProPainter inpainting).

Usage:
  vg-remove CLIP.mp4 --box 1700,40,1900,140        static region (logo/watermark bug)
  vg-remove CLIP.mp4 --point 640,360               moving object: SAM2 tracks it, then inpaint
  vg-remove CLIP.mp4 --mask mask.png               your own mask (white = remove);
                                                   also takes a dir of per-frame masks

Options: --out DIR (default CLIP_removed/), --frame N (SAM2 prompt frame for --point),
         --resize 0.5 (inpaint at half res — much faster/lighter, result upscaled back),
         --dilate N (mask dilation, default 4), --sub-len N (subvideo chunk, default 80)

Output in DIR: CLIP_clean.mov (ProRes LT, original audio) + CLIP_clean_preview.mp4
"""
import argparse, os, subprocess, sys
from pathlib import Path

os.environ["PATH"] = "/opt/homebrew/bin:" + os.environ.get("PATH", "")
os.environ.setdefault("PYTORCH_ENABLE_MPS_FALLBACK", "1")

VG = Path(__file__).resolve().parent.parent
PROPAINTER = VG / "tools/ProPainter"
RPY = VG / "venvs/roto/bin/python"
SAM2_CKPT = VG / "models/sam2.1_hiera_large.pt"
SAM2_CFG = "configs/sam2.1/sam2.1_hiera_l.yaml"


def run(cmd, **kw):
    r = subprocess.run([str(c) for c in cmd], **kw)
    if r.returncode != 0:
        sys.exit(f"command failed: {' '.join(str(c) for c in cmd)}")
    return r


def probe_video(path):
    out = subprocess.run(
        ["ffprobe", "-v", "error", "-select_streams", "v:0",
         "-show_entries", "stream=width,height,avg_frame_rate",
         "-of", "csv=p=0", str(path)], capture_output=True, text=True).stdout.strip()
    w, h, fr = out.split(",")[:3]
    num, den = fr.split("/")
    return int(w), int(h), (float(num) / float(den) if float(den) else 25.0)


def box_mask(w, h, box, dest):
    import numpy as np
    from PIL import Image
    x1, y1, x2, y2 = (int(v) for v in box)
    m = np.zeros((h, w), dtype=np.uint8)
    m[max(0, y1):min(h, y2), max(0, x1):min(w, x2)] = 255
    Image.fromarray(m).save(dest)


def sam2_track_masks(video, points, frame_n, masks_dir):
    import numpy as np, torch
    from PIL import Image
    from sam2.build_sam import build_sam2_video_predictor

    frames = masks_dir.parent / "frames"
    frames.mkdir(parents=True, exist_ok=True)
    if not list(frames.glob("*.jpg")):
        run(["ffmpeg", "-hide_banner", "-loglevel", "error", "-y", "-i", video,
             "-q:v", "2", "-start_number", "0", frames / "%05d.jpg"])
    device = "mps" if torch.backends.mps.is_available() else "cpu"
    print(f"SAM2: tracking object on {device} ...")
    pred = build_sam2_video_predictor(SAM2_CFG, str(SAM2_CKPT), device=device)
    state = pred.init_state(video_path=str(frames), offload_video_to_cpu=True,
                            offload_state_to_cpu=True)
    pc = np.array(points, dtype=np.float32)
    pl = np.ones(len(points), dtype=np.int32)
    pred.add_new_points_or_box(state, frame_idx=frame_n, obj_id=1, points=pc, labels=pl)
    masks_dir.mkdir(exist_ok=True)
    n = 0
    for fidx, _, logits in pred.propagate_in_video(state):
        m = (logits[0] > 0).cpu().numpy().squeeze().astype(np.uint8) * 255
        Image.fromarray(m).save(masks_dir / f"{fidx:05d}.png")
        n += 1
    print(f"SAM2: {n} frames tracked")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("video")
    ap.add_argument("--box", help="x1,y1,x2,y2 static region to remove")
    ap.add_argument("--point", action="append", default=[], help="x,y on the object (repeatable)")
    ap.add_argument("--mask", help="mask PNG (white = remove) or dir of per-frame masks")
    ap.add_argument("--frame", type=int, default=0)
    ap.add_argument("--resize", type=float, default=1.0, help="inpaint at this scale (0.5 = half res)")
    ap.add_argument("--dilate", type=int, default=4)
    ap.add_argument("--sub-len", type=int, default=80)
    ap.add_argument("--out", default=None)
    args = ap.parse_args()

    video = Path(args.video).expanduser().resolve()
    if not video.exists():
        sys.exit(f"no such file: {video}")
    stem = video.stem
    outdir = Path(args.out).expanduser() if args.out else video.parent / f"{stem}_removed"
    work = outdir / "work"
    work.mkdir(parents=True, exist_ok=True)
    w, h, fps = probe_video(video)

    if args.mask:
        mask = Path(args.mask).expanduser().resolve()
    elif args.box:
        mask = work / "box_mask.png"
        box_mask(w, h, [float(v) for v in args.box.split(",")], mask)
        print(f"static mask for box {args.box}")
    elif args.point:
        mask = work / "masks"
        points = [tuple(float(v) for v in p.split(",")) for p in args.point]
        sam2_track_masks(video, points, args.frame, mask)
    else:
        sys.exit("need --box, --point, or --mask")

    print("ProPainter: inpainting (the slow part) ...")
    cmd = [RPY, "inference_propainter.py", "--video", str(video), "--mask", str(mask),
           "--output", str(work), "--mask_dilation", str(args.dilate),
           "--subvideo_length", str(args.sub_len), "--save_fps", str(round(fps))]
    if args.resize != 1.0:
        cmd += ["--resize_ratio", str(args.resize)]
    r = subprocess.run([str(c) for c in cmd], cwd=PROPAINTER)
    if r.returncode != 0:
        sys.exit("ProPainter failed")

    results = sorted((work / stem).glob("inpaint_out.mp4")) or sorted(work.rglob("inpaint_out.mp4"))
    if not results:
        sys.exit(f"no inpaint_out.mp4 under {work}")
    clean = results[-1]

    out_mov = outdir / f"{stem}_clean.mov"
    out_prev = outdir / f"{stem}_clean_preview.mp4"
    print("composing ProRes + preview ...")
    # always scale back to source dims: imageio pads ProPainter's output to a
    # macro-block multiple (e.g. 804 -> 816) even at native res
    scale = f"scale={w}:{h}:flags=bicubic"
    run(["ffmpeg", "-hide_banner", "-loglevel", "error", "-y", "-i", clean, "-i", video,
         "-filter_complex", f"[0:v]{scale},format=yuv422p10le[v]",
         "-map", "[v]", "-map", "1:a:0?", "-c:v", "prores_ks", "-profile:v", "1",
         "-c:a", "pcm_s16le", "-shortest", out_mov])
    run(["ffmpeg", "-hide_banner", "-loglevel", "error", "-y", "-i", out_mov,
         "-c:v", "libx264", "-crf", "18", "-pix_fmt", "yuv420p", "-c:a", "aac", out_prev])
    print(f"\ndone:\n  {out_mov}\n  {out_prev}   <- preview")


if __name__ == "__main__":
    main()