commit 874aaaefa32029faeb672aad5510736d60c5842e Author: m3ultra Date: Sun Jul 26 22:02:20 2026 +1000 feat: ROTOGOD v0 — SAM-2 rotoscoping web app (load/scrub/trim → click actor → track → alpha export) FastAPI + PyAV + transformers SAM2.1 on MPS. Relative-URL frontend (works bare or behind digalot.fyi/rotogod/), ROTOGOD_ROOTS clip-path allowlist, fwd+rev propagation, ProRes 4444 / WebM / matte / greenscreen exports. deploy/rotogod.launchd.plist = always-on tailnet-bound instance. Co-Authored-By: Claude Fable 5 diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..7329662 --- /dev/null +++ b/.gitignore @@ -0,0 +1,6 @@ +.venv/ +__pycache__/ +*.pyc +output/ +samples/ +.DS_Store diff --git a/DESIGN.md b/DESIGN.md new file mode 100644 index 0000000..1a4b386 --- /dev/null +++ b/DESIGN.md @@ -0,0 +1,68 @@ +# ROTOGOD — design + +Local rotoscoping tool for m3ultra. Load a video → scrub → cut out a scene → click actors → +SAM 2 propagates masks across the clip → export alpha (ProRes 4444 / WebM / matte / greenscreen). + +## Why these choices + +| Decision | Choice | Why | +|---|---|---| +| UI | Local web app (FastAPI + vanilla JS canvas), port **8484** | Matches the fleet pattern (fluxgod :8890, paradramorama :8474, m3panel :8790). 8474–8479 are already taken on this box. Browser canvas is plenty fast for 720–1080p preview; no Qt dependency. | +| Segmentation | **SAM 2.1 via `transformers`** (`Sam2VideoModel` / `Sam2VideoProcessor`) | pip-installable (no facebookresearch git clone, no hydra configs), weights from HF hub, actively maintained, pure-torch so MPS works. API verified against installed transformers 5.14. | +| Model size | `facebook/sam2.1-hiera-small` default, `-large` for finals | Small = snappy interactive loop; large = better masks. Both cached in `~/.cache/huggingface`. Switch with `ROTOGOD_SAM2=facebook/sam2.1-hiera-large`. | +| Python | 3.12 venv via uv (`.venv/`) | Homebrew python3 is 3.14 — no torch wheels yet. | +| Decode | PyAV (bundled ffmpeg), seek-to-keyframe + roll forward, RAM frame cache | Arbitrary scrubbing; 256GB UMA means the cache can be generous (default 600 frames ≈ 1.6GB at 1080p). | +| Precision | float32 on MPS | bf16 on MPS still has op gaps; fp32 is safe and fast enough on M3 Ultra. | + +## Architecture + +``` +web/ (canvas UI: scrub, in/out, click prompts, overlays, export buttons) + │ JSON/JPEG/PNG over localhost:8484 +rotogod/server.py FastAPI — endpoints below, serves web/ +rotogod/video.py Clip: PyAV decode, frame(i) → RGB ndarray, JPEG previews, thumbstrip +rotogod/segment.py Sam2Engine (lazy model load) + RotoSession (frames, clicks, masks) +rotogod/jobs.py background threads with progress (propagate, exports) +rotogod/export.py ffmpeg: scene trim, ProRes 4444 alpha, VP9 alpha, matte, greenscreen +``` + +State is deliberately single-user/single-clip (module-level), same as fluxgod. Binds 127.0.0.1. + +## The loop + +1. **Open** a clip by path (`POST /api/open`). Meta + thumbstrip render; scrub via `GET /api/frame/{i}` (JPEG, cached decode). +2. **Cut the scene**: set in/out points (i/o keys), optional lossless keyframe-snapped trim or frame-exact re-encode (`POST /api/trim`). +3. **Roto session** (`POST /api/roto/start`): frames [in..out] are decoded once into RAM and handed to `processor.init_video_session(video=frames, inference_device="mps", video_storage_device="cpu")`. +4. **Click prompts**: left-click = positive, alt/right-click = negative, per object. The UI keeps the full point list per (frame, object) and sends it whole each time (SAM 2 replaces rather than accumulates per-frame inputs). Backend: `add_inputs_to_inference_session(...)` then `model(inference_session, frame_idx)` → instant mask overlay on that frame. +5. **Track** (`POST /api/roto/propagate`): background job runs `propagate_in_video_iterator` twice — forward from the prompted frame, then reverse to cover frames before it (SAM 2 only propagates one direction per pass); masks land per frame (packbits-compressed bools keyed by object) and the UI progress bar + overlay refresh as they stream in. +6. **Correct drift**: scrub to the bad frame, add +/- clicks there, re-track from that frame. SAM 2's memory bank handles re-propagation. +7. **Export** (`POST /api/export`): union of object masks → RGBA PNG sequence (optional edge feather) → ffmpeg: + - `prores`: ProRes 4444 `.mov`, `-pix_fmt yuva444p10le` (Resolve/FCP/Premiere-ready) + - `webm`: VP9 `yuva420p` (browser alpha) + - `matte`: grayscale h264 mp4 + - `green`: comp over green, h264 mp4 (quick eyeball check) + +## API surface + +``` +POST /api/open {path} GET /api/meta GET /api/frame/{i}?w= +GET /api/thumbs?n= POST /api/trim {in_f,out_f,exact} +POST /api/roto/start {in_f,out_f} POST /api/roto/click {frame,obj,points,labels} +POST /api/roto/propagate {start?} GET /api/roto/overlay/{frame}.png +GET /api/roto/status POST /api/roto/reset +POST /api/export {kind,feather} GET /api/job/{id} GET /api/status +``` + +Frame indices in the roto API are absolute clip frames; the session maps to local range internally. +Outputs land in `output/` (gitignored). + +## Known limits / phase 2+ + +- **Hair & motion blur**: SAM 2 masks are binary. The feather param is a stopgap; real fix is a + matting pass — **MatAnyone** (preferred) or RVM taking SAM 2's mask as guidance → soft alpha. + Slots in as a second engine in `segment.py` + an export-time toggle. +- **Long clips**: propagate memory grows with range; UI warns > 2000 frames. Chunk-and-stitch later. +- **Audio** is dropped in roto exports (kept in scene trims). +- **Multi-clip / project files**: out of scope for v0; one clip + one session at a time. +- Possible future: farm propagation to MODELBEAST queue (:8777) as a job type; not needed while + this box is the only one running it. diff --git a/README.md b/README.md new file mode 100644 index 0000000..d56c71d --- /dev/null +++ b/README.md @@ -0,0 +1,22 @@ +# ROTOGOD + +Click an actor, get an alpha channel. Local SAM 2 rotoscoping on Apple Silicon. + +```bash +# setup (once) +uv venv -p 3.12 .venv +uv pip install --python .venv/bin/python -r requirements.txt + +# run +.venv/bin/python run.py --port 8484 # → http://localhost:8484 +``` + +- Load a clip by absolute path, scrub with the slider / arrow keys / mouse wheel. +- `i` / `o` set in/out points. Trim exports a scene cut (`output/`). +- **Start roto** locks the in/out range into a SAM 2 session. Left-click = include, + right-click = exclude. Multiple objects supported. **Track** propagates through the range. +- Fix drift by clicking on the bad frame, then Track again from there. +- Export ProRes 4444 (alpha), WebM (alpha), matte, or greenscreen. + +Model defaults to `facebook/sam2.1-hiera-small`; for final-quality masks run with +`ROTOGOD_SAM2=facebook/sam2.1-hiera-large`. See [DESIGN.md](DESIGN.md). diff --git a/deploy/rotogod.launchd.plist b/deploy/rotogod.launchd.plist new file mode 100644 index 0000000..5382690 --- /dev/null +++ b/deploy/rotogod.launchd.plist @@ -0,0 +1,27 @@ + + + + + + Labelparty.monster.rotogod + ProgramArguments + + /Users/m3ultra/Documents/rotogod/.venv/bin/python + /Users/m3ultra/Documents/rotogod/run.py + --host100.89.131.57 + --port8484 + + WorkingDirectory/Users/m3ultra/Documents/rotogod + RunAtLoad + KeepAlive + StandardOutPath/Users/m3ultra/Library/Logs/rotogod.log + StandardErrorPath/Users/m3ultra/Library/Logs/rotogod.log + + diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..11f1586 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,10 @@ +torch +torchvision +transformers +huggingface_hub +fastapi +uvicorn[standard] +av +opencv-python-headless +pillow +numpy diff --git a/rotogod/__init__.py b/rotogod/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/rotogod/export.py b/rotogod/export.py new file mode 100644 index 0000000..8cbd75c --- /dev/null +++ b/rotogod/export.py @@ -0,0 +1,88 @@ +"""ffmpeg exports: scene trims and alpha/matte renders. Outputs land in output/.""" + +import os +import shutil +import subprocess +import tempfile + +import cv2 +import numpy as np + +ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +OUTPUT_DIR = os.path.join(ROOT, "output") + +GREEN = np.array([0, 177, 64], np.uint8) # RGB + + +def _run(cmd: list[str]): + r = subprocess.run(cmd, capture_output=True, text=True) + if r.returncode != 0: + raise RuntimeError(f"ffmpeg failed: {r.stderr[-800:]}") + + +def trim(src: str, in_f: int, out_f: int, fps: float, exact: bool) -> str: + os.makedirs(OUTPUT_DIR, exist_ok=True) + stem = os.path.splitext(os.path.basename(src))[0] + t0, dur = in_f / fps, (out_f - in_f + 1) / fps + if exact: + dst = os.path.join(OUTPUT_DIR, f"{stem}_cut_{in_f}-{out_f}.mp4") + cmd = ["ffmpeg", "-y", "-ss", f"{t0:.6f}", "-i", src, "-t", f"{dur:.6f}", + "-c:v", "libx264", "-crf", "16", "-preset", "fast", + "-pix_fmt", "yuv420p", "-c:a", "aac", dst] + else: + dst = os.path.join(OUTPUT_DIR, f"{stem}_cut_{in_f}-{out_f}_copy.mp4") + cmd = ["ffmpeg", "-y", "-ss", f"{t0:.6f}", "-i", src, "-t", f"{dur:.6f}", + "-c", "copy", dst] + _run(cmd) + return dst + + +def export_roto(clip, sess, kind: str = "prores", feather: int = 0, + progress_cb=None) -> str: + """Render the roto range with the union of object masks as alpha.""" + os.makedirs(OUTPUT_DIR, exist_ok=True) + work = tempfile.mkdtemp(prefix="work_", dir=OUTPUT_DIR) + obj_ids = sess.object_ids() + try: + for li in range(sess.n): + rgb = clip.frame(sess.in_f + li) + alpha = np.zeros((sess.height, sess.width), np.uint8) + for oid in obj_ids: + m = sess.get_mask(oid, li) + if m is not None: + alpha[m] = 255 + if feather > 0: + k = feather * 2 + 1 + alpha = cv2.GaussianBlur(alpha, (k, k), 0) + p = os.path.join(work, f"{li:06d}.png") + if kind == "matte": + cv2.imwrite(p, alpha) + elif kind == "green": + a = alpha.astype(np.float32)[..., None] / 255.0 + comp = (rgb * a + GREEN * (1 - a)).astype(np.uint8) + cv2.imwrite(p, cv2.cvtColor(comp, cv2.COLOR_RGB2BGR)) + else: # prores / webm: RGBA + bgra = np.dstack([cv2.cvtColor(rgb, cv2.COLOR_RGB2BGR), alpha]) + cv2.imwrite(p, bgra) + if progress_cb: + progress_cb(li + 1, sess.n) + + stem = os.path.splitext(os.path.basename(clip.path))[0] + base = os.path.join(OUTPUT_DIR, f"{stem}_roto_{sess.in_f}-{sess.out_f}") + seq = ["-framerate", f"{clip.fps}", "-i", os.path.join(work, "%06d.png")] + if kind == "prores": + dst = base + ".mov" + cmd = ["ffmpeg", "-y", *seq, "-c:v", "prores_ks", "-profile:v", "4444", + "-pix_fmt", "yuva444p10le", "-vendor", "apl0", dst] + elif kind == "webm": + dst = base + ".webm" + cmd = ["ffmpeg", "-y", *seq, "-c:v", "libvpx-vp9", "-pix_fmt", "yuva420p", + "-auto-alt-ref", "0", "-b:v", "0", "-crf", "24", dst] + else: # matte / green -> h264 + dst = base + f"_{kind}.mp4" + cmd = ["ffmpeg", "-y", *seq, "-c:v", "libx264", "-crf", "16", + "-pix_fmt", "yuv420p", dst] + _run(cmd) + return dst + finally: + shutil.rmtree(work, ignore_errors=True) diff --git a/rotogod/jobs.py b/rotogod/jobs.py new file mode 100644 index 0000000..b9dc0c4 --- /dev/null +++ b/rotogod/jobs.py @@ -0,0 +1,44 @@ +"""Background jobs with progress, polled by the UI.""" + +import threading +import time +import traceback +import uuid + +JOBS: dict[str, "Job"] = {} + + +class Job: + def __init__(self, kind: str): + self.id = uuid.uuid4().hex[:8] + self.kind = kind + self.status = "running" + self.progress = 0.0 + self.message = "" + self.result = None + self.error = None + self.t0 = time.time() + + def to_dict(self): + return {"id": self.id, "kind": self.kind, "status": self.status, + "progress": round(self.progress, 4), "message": self.message, + "result": self.result, "error": self.error, + "elapsed": round(time.time() - self.t0, 1)} + + +def start(kind: str, fn) -> Job: + """fn(job) runs in a thread; sets job.result on success.""" + job = Job(kind) + JOBS[job.id] = job + + def run(): + try: + job.result = fn(job) + job.status = "done" + job.progress = 1.0 + except Exception: + job.status = "error" + job.error = traceback.format_exc(limit=8) + + threading.Thread(target=run, daemon=True).start() + return job diff --git a/rotogod/segment.py b/rotogod/segment.py new file mode 100644 index 0000000..10d91b4 --- /dev/null +++ b/rotogod/segment.py @@ -0,0 +1,144 @@ +"""SAM 2.1 video segmentation via transformers, on MPS.""" + +import os +import threading + +import numpy as np +import torch + +MODEL_ID = os.environ.get("ROTOGOD_SAM2", "facebook/sam2.1-hiera-small") + + +class Sam2Engine: + def __init__(self): + self.device = "mps" if torch.backends.mps.is_available() else "cpu" + self.dtype = torch.float32 + self.model = None + self.processor = None + self.status = "cold" + self._load_lock = threading.Lock() + + def ensure_loaded(self): + with self._load_lock: + if self.model is not None: + return + self.status = "loading" + try: + from transformers import Sam2VideoModel, Sam2VideoProcessor + self.processor = Sam2VideoProcessor.from_pretrained(MODEL_ID) + self.model = (Sam2VideoModel.from_pretrained(MODEL_ID, dtype=self.dtype) + .to(self.device).eval()) + self.status = f"ready · {MODEL_ID.split('/')[-1]} · {self.device}" + except Exception as e: + self.status = f"error: {e}" + self.model = self.processor = None + raise + + +ENGINE = Sam2Engine() + + +def _ids(obj_ids) -> list[int]: + if obj_ids is None: + return [] + return obj_ids.tolist() if hasattr(obj_ids, "tolist") else list(obj_ids) + + +class RotoSession: + """One SAM2 inference session over clip frames [in_f..out_f] inclusive. + + Frame indices in the public API are absolute clip frames; internally local. + Masks are stored packbits-compressed per (obj_id, local_idx). + """ + + def __init__(self, frames: list[np.ndarray], in_f: int, out_f: int, + width: int, height: int): + ENGINE.ensure_loaded() + self.in_f, self.out_f = in_f, out_f + self.width, self.height = width, height + self.n = len(frames) + self.model_lock = threading.Lock() # SAM2 session is not reentrant + self.session = ENGINE.processor.init_video_session( + video=frames, + inference_device=ENGINE.device, + video_storage_device="cpu", + dtype=ENGINE.dtype, + ) + self.masks: dict[tuple[int, int], np.ndarray] = {} + self.clicks: dict[tuple[int, int], dict] = {} # (obj, local) -> {points, labels} + self.propagating = False + self.tracked_frames = 0 + + # -- helpers ----------------------------------------------------------- + + def local(self, abs_frame: int) -> int: + li = abs_frame - self.in_f + if not 0 <= li < self.n: + raise ValueError(f"frame {abs_frame} outside roto range {self.in_f}..{self.out_f}") + return li + + def object_ids(self) -> list[int]: + ids = {k[0] for k in self.masks} | {k[0] for k in self.clicks} + return sorted(ids) + + def get_mask(self, obj_id: int, local_idx: int) -> np.ndarray | None: + p = self.masks.get((obj_id, local_idx)) + if p is None: + return None + return (np.unpackbits(p, count=self.height * self.width) + .reshape(self.height, self.width).astype(bool)) + + def _store(self, obj_id: int, local_idx: int, mask: np.ndarray): + self.masks[(obj_id, local_idx)] = np.packbits(mask) + + def _postprocess(self, pred_masks) -> np.ndarray: + m = ENGINE.processor.post_process_masks( + [pred_masks], original_sizes=[[self.height, self.width]], binarize=True, + max_hole_area=200.0, max_sprinkle_area=200.0)[0] + if m.ndim == 4: + m = m[:, 0] + return m.cpu().numpy().astype(bool) # (num_obj, H, W) + + # -- inference --------------------------------------------------------- + + def click(self, abs_frame: int, obj_id: int, points: list[list[float]], + labels: list[int]): + """points/labels are the FULL prompt set for this frame+object + (SAM2 replaces per-frame inputs rather than accumulating).""" + li = self.local(abs_frame) + with self.model_lock: + self.clicks[(obj_id, li)] = {"points": points, "labels": labels} + ENGINE.processor.add_inputs_to_inference_session( + inference_session=self.session, + frame_idx=li, + obj_ids=obj_id, + input_points=[[points]], + input_labels=[[labels]], + ) + with torch.inference_mode(): + out = ENGINE.model(inference_session=self.session, frame_idx=li) + masks = self._postprocess(out.pred_masks) + for i, oid in enumerate(_ids(out.object_ids)): + self._store(oid, li, masks[i]) + + def propagate(self, start_abs: int | None = None, progress_cb=None): + """Forward pass from the prompted frame, then a reverse pass so the + frames before the prompt get covered too.""" + start_local = self.local(start_abs) if start_abs is not None else None + with self.model_lock: + self.propagating = True + self.tracked_frames = 0 + try: + with torch.inference_mode(): + for reverse in (False, True): + for out in ENGINE.model.propagate_in_video_iterator( + self.session, start_frame_idx=start_local, + reverse=reverse): + masks = self._postprocess(out.pred_masks) + for i, oid in enumerate(_ids(out.object_ids)): + self._store(oid, out.frame_idx, masks[i]) + self.tracked_frames += 1 + if progress_cb: + progress_cb(out.frame_idx, self.tracked_frames) + finally: + self.propagating = False diff --git a/rotogod/server.py b/rotogod/server.py new file mode 100644 index 0000000..86a1523 --- /dev/null +++ b/rotogod/server.py @@ -0,0 +1,279 @@ +"""rotogod server — FastAPI over one clip + one roto session (single-user tool).""" + +import contextlib +import os +import threading + +import cv2 +import numpy as np +from fastapi import FastAPI, HTTPException +from fastapi.responses import FileResponse, Response +from pydantic import BaseModel + +from . import export, jobs, segment +from .video import Clip + +ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +WEB = os.path.join(ROOT, "web") + +# object colors, RGB — keep in sync with app.js PALETTE +PALETTE = [(122, 208, 255), (255, 138, 101), (129, 255, 141), + (255, 121, 198), (255, 216, 107), (167, 139, 250)] + +STATE: dict = {"clip": None, "session": None} + + +@contextlib.asynccontextmanager +async def lifespan(app): + if os.environ.get("ROTOGOD_PRELOAD", "1") == "1": + threading.Thread(target=lambda: _safe_load(), daemon=True).start() + yield + + +def _safe_load(): + try: + segment.ENGINE.ensure_loaded() + except Exception: + pass # status carries the error; UI shows it + + +app = FastAPI(title="rotogod", lifespan=lifespan) + + +def clip() -> Clip: + c = STATE["clip"] + if c is None: + raise HTTPException(400, "no clip loaded") + return c + + +def sess() -> segment.RotoSession: + s = STATE["session"] + if s is None: + raise HTTPException(400, "no roto session — set in/out and Start roto") + return s + + +# -- static --------------------------------------------------------------- + +@app.get("/") +def index(): + return FileResponse(os.path.join(WEB, "index.html")) + + +@app.get("/static/{name}") +def static(name: str): + p = os.path.join(WEB, os.path.basename(name)) + if not os.path.isfile(p): + raise HTTPException(404) + return FileResponse(p) + + +# -- clip ----------------------------------------------------------------- + +class OpenBody(BaseModel): + path: str + + +# clip paths must live under one of these roots (ship-check #3: no open path oracle) +ROOTS = [os.path.realpath(os.path.expanduser(p)) + for p in os.environ.get("ROTOGOD_ROOTS", "~").split(":") if p.strip()] + + +@app.post("/api/open") +def open_clip(body: OpenBody): + path = os.path.realpath(os.path.expanduser(body.path.strip())) + if not any(path == r or path.startswith(r + os.sep) for r in ROOTS): + raise HTTPException(403, "path outside ROTOGOD_ROOTS") + if not os.path.isfile(path): + raise HTTPException(400, f"not a file: {path}") + old = STATE["clip"] + try: + STATE["clip"] = Clip(path) + except Exception as e: + raise HTTPException(400, f"could not open: {e}") + STATE["session"] = None + if old: + old.close() + return meta() + + +@app.get("/api/meta") +def meta(): + c = clip() + return {"path": c.path, "width": c.width, "height": c.height, + "fps": c.fps, "nframes": c.nframes, "duration": c.duration} + + +@app.get("/api/frame/{idx}") +def frame(idx: int, w: int = 1280): + return Response(clip().jpeg(idx, maxw=w), media_type="image/jpeg") + + +@app.get("/api/thumbs") +def thumbs(n: int = 60): + return Response(clip().thumb_sprite(n=n), media_type="image/jpeg") + + +class TrimBody(BaseModel): + in_f: int + out_f: int + exact: bool = True + + +@app.post("/api/trim") +def trim(body: TrimBody): + c = clip() + job = jobs.start("trim", lambda j: export.trim( + c.path, body.in_f, body.out_f, c.fps, body.exact)) + return {"job": job.id} + + +# -- roto ----------------------------------------------------------------- + +class RotoStartBody(BaseModel): + in_f: int + out_f: int + + +@app.post("/api/roto/start") +def roto_start(body: RotoStartBody): + c = clip() + in_f = max(0, body.in_f) + out_f = min(c.nframes - 1, body.out_f) + if out_f <= in_f: + raise HTTPException(400, "out point must be after in point") + if out_f - in_f + 1 > 2000: + raise HTTPException(400, "range > 2000 frames — trim the scene first") + + def build(job): + job.message = "decoding frames" + frames = [] + for i in range(in_f, out_f + 1): + frames.append(c.frame(i).copy()) + job.progress = 0.5 * (i - in_f + 1) / (out_f - in_f + 1) + job.message = "building SAM2 session" + STATE["session"] = segment.RotoSession(frames, in_f, out_f, c.width, c.height) + job.message = "ready" + return {"in_f": in_f, "out_f": out_f, "n": out_f - in_f + 1} + + return {"job": jobs.start("roto_start", build).id} + + +class ClickBody(BaseModel): + frame: int + obj: int + points: list[list[float]] # full set for this frame+obj, video-pixel coords + labels: list[int] # 1 = include, 0 = exclude + + +@app.post("/api/roto/click") +def roto_click(body: ClickBody): + s = sess() + if s.propagating: + raise HTTPException(409, "propagation running — wait for it to finish") + try: + s.click(body.frame, body.obj, body.points, body.labels) + except ValueError as e: + raise HTTPException(400, str(e)) + return {"ok": True} + + +class PropagateBody(BaseModel): + start: int | None = None # absolute frame; None = from first prompted frame + + +@app.post("/api/roto/propagate") +def roto_propagate(body: PropagateBody): + s = sess() + if s.propagating: + raise HTTPException(409, "already propagating") + if not s.clicks: + raise HTTPException(400, "no prompts — click the actor first") + + def run(job): + def cb(frame_idx, done): + job.progress = min(1.0, done / s.n) + job.message = f"frame {s.in_f + frame_idx}" + s.propagate(start_abs=body.start, progress_cb=cb) + return {"tracked": s.tracked_frames} + + return {"job": jobs.start("propagate", run).id} + + +@app.get("/api/roto/overlay/{frame_idx}.png") +def roto_overlay(frame_idx: int): + s = sess() + try: + li = s.local(frame_idx) + except ValueError as e: + raise HTTPException(400, str(e)) + rgba = np.zeros((s.height, s.width, 4), np.uint8) + for oid in s.object_ids(): + m = s.get_mask(oid, li) + if m is None: + continue + r, g, b = PALETTE[(oid - 1) % len(PALETTE)] + rgba[m] = (b, g, r, 110) + edges = cv2.findContours(m.astype(np.uint8), cv2.RETR_EXTERNAL, + cv2.CHAIN_APPROX_SIMPLE)[0] + cv2.drawContours(rgba, edges, -1, (b, g, r, 255), 2) + ok, buf = cv2.imencode(".png", rgba) + return Response(buf.tobytes(), media_type="image/png") + + +@app.get("/api/roto/status") +def roto_status(): + s = STATE["session"] + if s is None: + return {"active": False} + return {"active": True, "in_f": s.in_f, "out_f": s.out_f, "n": s.n, + "objects": s.object_ids(), "propagating": s.propagating, + "masked_frames": len({k[1] for k in s.masks})} + + +@app.post("/api/roto/reset") +def roto_reset(): + STATE["session"] = None + return {"ok": True} + + +class ExportBody(BaseModel): + kind: str = "prores" # prores | webm | matte | green + feather: int = 0 + + +@app.post("/api/export") +def export_roto(body: ExportBody): + c, s = clip(), sess() + if not s.masks: + raise HTTPException(400, "no masks yet — Track first") + + def run(job): + def cb(done, total): + job.progress = 0.8 * done / total + job.message = f"rendering {done}/{total}" + job.message = "encoding" + return export.export_roto(c, s, kind=body.kind, feather=body.feather, + progress_cb=cb) + + return {"job": jobs.start(f"export_{body.kind}", run).id} + + +# -- misc ----------------------------------------------------------------- + +@app.get("/api/job/{job_id}") +def job(job_id: str): + j = jobs.JOBS.get(job_id) + if j is None: + raise HTTPException(404) + return j.to_dict() + + +@app.get("/api/status") +def status(): + c = STATE["clip"] + return {"engine": segment.ENGINE.status, + "model": segment.MODEL_ID, + "clip": meta() if c else None, + "roto": roto_status()} diff --git a/rotogod/video.py b/rotogod/video.py new file mode 100644 index 0000000..4494b19 --- /dev/null +++ b/rotogod/video.py @@ -0,0 +1,115 @@ +"""PyAV-backed clip with arbitrary frame seeking and a RAM frame cache.""" + +import threading + +import av +import cv2 +import numpy as np + + +class Clip: + CACHE_CAP = 600 # decoded frames kept in RAM (~1.6GB at 1080p) + ROLL_AHEAD = 48 # decode forward without seeking if target is this close + + def __init__(self, path: str): + self.path = path + self.lock = threading.RLock() + self.container = av.open(path) + self.stream = self.container.streams.video[0] + self.stream.thread_type = "AUTO" + cc = self.stream.codec_context + self.width, self.height = cc.width, cc.height + self.fps = float(self.stream.average_rate or self.stream.guessed_rate or 25) + self.time_base = self.stream.time_base + self.start_pts = self.stream.start_time or 0 + + n = self.stream.frames or 0 + if not n: + if self.stream.duration: + dur = float(self.stream.duration * self.time_base) + elif self.container.duration: + dur = self.container.duration / av.time_base + else: + dur = 0 + n = int(round(dur * self.fps)) + self.nframes = max(int(n), 1) + self.duration = self.nframes / self.fps + + self._cache: dict[int, np.ndarray] = {} + self._order: list[int] = [] + self._iter = None + self._last = -10 ** 9 + self._thumbs: bytes | None = None + + def close(self): + with self.lock: + self.container.close() + + # -- decoding ---------------------------------------------------------- + + def _pts_to_idx(self, pts) -> int: + if pts is None: + return self._last + 1 + return int(round(float((pts - self.start_pts) * self.time_base) * self.fps)) + + def _store(self, idx: int, arr: np.ndarray): + if idx in self._cache: + return + self._cache[idx] = arr + self._order.append(idx) + while len(self._order) > self.CACHE_CAP: + self._cache.pop(self._order.pop(0), None) + + def frame(self, idx: int) -> np.ndarray: + """Decoded RGB frame at index (clamped).""" + idx = max(0, min(idx, self.nframes - 1)) + with self.lock: + hit = self._cache.get(idx) + if hit is not None: + return hit + if self._iter is None or not (0 <= idx - self._last <= self.ROLL_AHEAD): + ts = self.start_pts + int(idx / self.fps / self.time_base) + self.container.seek(ts, stream=self.stream, backward=True) + self._iter = self.container.decode(self.stream) + self._last = -10 ** 9 + last_arr = None + for frame in self._iter: + fi = self._pts_to_idx(frame.pts) + self._last = fi + arr = frame.to_ndarray(format="rgb24") + self._store(fi, arr) + last_arr = arr + if fi >= idx: + return arr + # hit EOF before reaching idx — metadata overstated the count + self._iter = None + if last_arr is not None: + self.nframes = min(self.nframes, self._last + 1) + return last_arr + raise RuntimeError(f"failed to decode frame {idx} of {self.path}") + + # -- previews ---------------------------------------------------------- + + def jpeg(self, idx: int, maxw: int = 1280, quality: int = 87) -> bytes: + rgb = self.frame(idx) + if maxw and rgb.shape[1] > maxw: + h = int(rgb.shape[0] * maxw / rgb.shape[1]) + rgb = cv2.resize(rgb, (maxw, h), interpolation=cv2.INTER_AREA) + ok, buf = cv2.imencode(".jpg", cv2.cvtColor(rgb, cv2.COLOR_RGB2BGR), + [cv2.IMWRITE_JPEG_QUALITY, quality]) + return buf.tobytes() + + def thumb_sprite(self, n: int = 60, w: int = 120) -> bytes: + """Horizontal strip of n evenly spaced thumbnails, cached after first build.""" + if self._thumbs is not None: + return self._thumbs + h = max(1, int(w * self.height / self.width)) + strip = np.zeros((h, w * n, 3), np.uint8) + for k in range(n): + idx = int(k * (self.nframes - 1) / max(n - 1, 1)) + rgb = self.frame(idx) + strip[:, k * w:(k + 1) * w] = cv2.resize(rgb, (w, h), interpolation=cv2.INTER_AREA) + ok, buf = cv2.imencode(".jpg", cv2.cvtColor(strip, cv2.COLOR_RGB2BGR), + [cv2.IMWRITE_JPEG_QUALITY, 70]) + self._thumbs = buf.tobytes() + return self._thumbs diff --git a/run.py b/run.py new file mode 100644 index 0000000..5ab103d --- /dev/null +++ b/run.py @@ -0,0 +1,16 @@ +import argparse +import os +import sys + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) + +import uvicorn + +from rotogod.server import app + +if __name__ == "__main__": + p = argparse.ArgumentParser() + p.add_argument("--port", type=int, default=8484) + p.add_argument("--host", default="127.0.0.1") + a = p.parse_args() + uvicorn.run(app, host=a.host, port=a.port, log_level="warning") diff --git a/web/app.js b/web/app.js new file mode 100644 index 0000000..dfec8eb --- /dev/null +++ b/web/app.js @@ -0,0 +1,307 @@ +const $ = s => document.querySelector(s); +const sleep = ms => new Promise(r => setTimeout(r, ms)); +const PALETTE = ['#7ad0ff', '#ff8a65', '#81ff8d', '#ff79c6', '#ffd86b', '#a78bfa']; + +const cv = $('#cv'), ctx = cv.getContext('2d'); +let meta = null, cur = 0, inF = 0, outF = 0; +let roto = { active: false }; +let objects = [1], activeObj = 1; +let clicks = {}; // obj -> frame -> {points:[[x,y]..], labels:[..]} +let playTimer = null; +let frameToken = 0, overlayV = 0; + +// ---------- api / ui plumbing ---------- + +async function api(path, body) { + const r = await fetch(path, body ? { + method: 'POST', headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify(body) + } : undefined); + if (!r.ok) { + const d = await r.json().catch(() => ({})); + throw new Error(d.detail || r.statusText); + } + return r.json(); +} + +let toastTimer = null; +function toast(msg, err) { + const t = $('#toast'); + t.textContent = msg; + t.className = err ? 'err' : ''; + t.style.opacity = 1; + clearTimeout(toastTimer); + toastTimer = setTimeout(() => t.style.opacity = 0, err ? 6000 : 3000); +} + +function loadImg(src) { + return new Promise((res, rej) => { + const im = new Image(); + im.onload = () => res(im); + im.onerror = rej; + im.src = src; + }); +} + +async function pollJob(id, onprog) { + for (;;) { + const j = await api(`api/job/${id}`); + $('#prog').value = j.progress; + $('#progText').textContent = j.status === 'running' ? (j.message || j.kind) : ''; + if (onprog) await onprog(j); + if (j.status === 'done') { $('#prog').value = 0; return j; } + if (j.status === 'error') { + $('#prog').value = 0; + throw new Error((j.error || 'job failed').trim().split('\n').pop()); + } + await sleep(300); + } +} + +// ---------- frame display ---------- + +async function showFrame(i) { + if (!meta) return; + cur = Math.max(0, Math.min(i, meta.nframes - 1)); + $('#scrub').value = cur; + updateLabels(); + const tok = ++frameToken; + try { + const img = await loadImg(`api/frame/${cur}?w=1280`); + if (tok !== frameToken) return; + ctx.drawImage(img, 0, 0, cv.width, cv.height); + if (roto.active && $('#showOverlay').checked && cur >= roto.in_f && cur <= roto.out_f) { + try { + const ov = await loadImg(`api/roto/overlay/${cur}.png?v=${overlayV}`); + if (tok !== frameToken) return; + ctx.drawImage(ov, 0, 0, cv.width, cv.height); + } catch (e) { /* no overlay yet */ } + drawClickMarkers(); + } + } catch (e) { /* frame fetch raced a reload */ } +} + +function drawClickMarkers() { + for (const obj of objects) { + const c = (clicks[obj] || {})[cur]; + if (!c) continue; + c.points.forEach(([x, y], k) => { + ctx.beginPath(); + ctx.arc(x, y, 7, 0, Math.PI * 2); + ctx.fillStyle = c.labels[k] ? '#28d478' : '#ff5252'; + ctx.fill(); + ctx.strokeStyle = '#fff'; + ctx.lineWidth = 2; + ctx.stroke(); + }); + } +} + +function updateLabels() { + if (!meta) return; + $('#frameLabel').textContent = + `${cur} / ${meta.nframes - 1} · ${(cur / meta.fps).toFixed(2)}s`; + $('#rangeInfo').textContent = roto.active + ? `session: ${roto.in_f} → ${roto.out_f}` + : `range: ${inF} → ${outF} (${outF - inF + 1}f)`; + const n = meta.nframes; + $('#shadeL').style.width = `${inF / n * 100}%`; + $('#shadeR').style.width = `${(n - 1 - outF) / n * 100}%`; + $('#playhead').style.left = `${cur / Math.max(n - 1, 1) * 100}%`; +} + +// ---------- clip ---------- + +async function loadClip() { + const path = $('#path').value.trim(); + if (!path) return toast('enter a clip path', true); + stopPlay(); + try { + meta = await api('api/open', { path }); + } catch (e) { return toast(e.message, true); } + localStorage.setItem('rotogod_path', path); + cur = 0; inF = 0; outF = meta.nframes - 1; + roto = { active: false }; clicks = {}; renderObjects(); + cv.width = meta.width; cv.height = meta.height; + $('#scrub').max = meta.nframes - 1; + $('#thumbs').src = `api/thumbs?n=60&v=${Date.now()}`; + await showFrame(0); + toast(`${meta.width}×${meta.height} · ${meta.fps.toFixed(2)}fps · ${meta.nframes}f`); +} + +// ---------- playback / scrub ---------- + +function stopPlay() { + if (playTimer) { clearInterval(playTimer); playTimer = null; $('#play').textContent = '▶'; } +} +function togglePlay() { + if (!meta) return; + if (playTimer) return stopPlay(); + $('#play').textContent = '⏸'; + const loopIn = outF > inF ? inF : 0, loopOut = outF > inF ? outF : meta.nframes - 1; + playTimer = setInterval(() => { + showFrame(cur >= loopOut ? loopIn : cur + 1); + }, 1000 / meta.fps); +} + +// ---------- roto ---------- + +function renderObjects() { + const box = $('#objects'); + box.innerHTML = ''; + for (const id of objects) { + const el = document.createElement('div'); + el.className = 'obj' + (id === activeObj ? ' active' : ''); + el.innerHTML = `obj ${id}`; + el.onclick = () => { activeObj = id; renderObjects(); }; + box.appendChild(el); + } +} + +async function startRoto() { + if (!meta) return toast('load a clip first', true); + stopPlay(); + try { + const { job } = await api('api/roto/start', { in_f: inF, out_f: outF }); + await pollJob(job); + roto = await api('api/roto/status'); + clicks = {}; overlayV++; + await showFrame(Math.max(roto.in_f, Math.min(cur, roto.out_f))); + toast('session ready — click the actor'); + } catch (e) { toast(e.message, true); } +} + +async function sendClick(x, y, label) { + clicks[activeObj] = clicks[activeObj] || {}; + const c = clicks[activeObj][cur] = clicks[activeObj][cur] || { points: [], labels: [] }; + c.points.push([Math.round(x), Math.round(y)]); + c.labels.push(label); + try { + await api('api/roto/click', { frame: cur, obj: activeObj, points: c.points, labels: c.labels }); + overlayV++; + await showFrame(cur); + } catch (e) { + c.points.pop(); c.labels.pop(); + toast(e.message, true); + } +} + +async function track(fromHere) { + stopPlay(); + try { + const { job } = await api('api/roto/propagate', { start: fromHere ? cur : null }); + let tick = 0; + await pollJob(job, async () => { + if (++tick % 2 === 0) { overlayV++; await showFrame(cur); } + }); + overlayV++; + await showFrame(cur); + toast('tracked — scrub to check, click bad frames to fix'); + } catch (e) { toast(e.message, true); } +} + +// ---------- exports ---------- + +function addOutput(p) { + const d = document.createElement('div'); + d.textContent = p; + $('#outlist').prepend(d); +} + +async function exportRoto(kind) { + try { + const { job } = await api('api/export', { kind, feather: +$('#feather').value }); + const j = await pollJob(job); + addOutput(j.result); + toast(`${kind} done`); + } catch (e) { toast(e.message, true); } +} + +async function exportTrim() { + if (!meta) return; + try { + const { job } = await api('api/trim', { in_f: inF, out_f: outF, exact: $('#exactTrim').checked }); + const j = await pollJob(job); + addOutput(j.result); + toast('trim done'); + } catch (e) { toast(e.message, true); } +} + +// ---------- events ---------- + +$('#load').onclick = loadClip; +$('#path').addEventListener('keydown', e => { if (e.key === 'Enter') loadClip(); }); +$('#play').onclick = togglePlay; +$('#prev').onclick = () => { stopPlay(); showFrame(cur - 1); }; +$('#next').onclick = () => { stopPlay(); showFrame(cur + 1); }; +$('#prev10').onclick = () => { stopPlay(); showFrame(cur - 10); }; +$('#next10').onclick = () => { stopPlay(); showFrame(cur + 10); }; +$('#setIn').onclick = () => { inF = Math.min(cur, outF); updateLabels(); }; +$('#setOut').onclick = () => { outF = Math.max(cur, inF); updateLabels(); }; +$('#scrub').addEventListener('input', e => { stopPlay(); showFrame(+e.target.value); }); +$('#startRoto').onclick = startRoto; +$('#resetRoto').onclick = async () => { + await api('api/roto/reset', {}); roto = { active: false }; clicks = {}; + updateLabels(); showFrame(cur); toast('session dropped'); +}; +$('#addObj').onclick = () => { + objects.push(objects.length + 1); + activeObj = objects.length; + renderObjects(); +}; +$('#track').onclick = () => track(false); +$('#trackHere').onclick = () => track(true); +$('#showOverlay').onchange = () => showFrame(cur); +$('#trim').onclick = exportTrim; +document.querySelectorAll('.exp').forEach(b => b.onclick = () => exportRoto(b.dataset.kind)); + +cv.addEventListener('pointerdown', e => { + if (!meta) return; + if (!roto.active) return toast('set in/out then Start roto'); + if (cur < roto.in_f || cur > roto.out_f) return toast('outside session range', true); + stopPlay(); + const r = cv.getBoundingClientRect(); + const x = (e.clientX - r.left) / r.width * meta.width; + const y = (e.clientY - r.top) / r.height * meta.height; + sendClick(x, y, (e.button === 2 || e.altKey) ? 0 : 1); +}); +cv.addEventListener('contextmenu', e => e.preventDefault()); + +$('#stage').addEventListener('wheel', e => { + if (!meta) return; + e.preventDefault(); + stopPlay(); + showFrame(cur + Math.sign(e.deltaY)); +}, { passive: false }); + +$('#timeline').addEventListener('pointerdown', e => { + if (!meta) return; + stopPlay(); + const r = e.currentTarget.getBoundingClientRect(); + showFrame(Math.round((e.clientX - r.left) / r.width * (meta.nframes - 1))); +}); + +document.addEventListener('keydown', e => { + if (!meta || /INPUT/.test(e.target.tagName)) return; + const step = e.shiftKey ? 10 : 1; + if (e.key === 'ArrowLeft') { stopPlay(); showFrame(cur - step); } + else if (e.key === 'ArrowRight') { stopPlay(); showFrame(cur + step); } + else if (e.key === 'i') { inF = Math.min(cur, outF); updateLabels(); } + else if (e.key === 'o') { outF = Math.max(cur, inF); updateLabels(); } + else if (e.key === ' ') { e.preventDefault(); togglePlay(); } +}); + +// ---------- boot ---------- + +renderObjects(); +$('#path').value = localStorage.getItem('rotogod_path') || ''; +(async function engineBadge() { + for (;;) { + try { + const s = await api('api/status'); + $('#engine').textContent = s.engine.startsWith('ready') ? s.engine : `engine: ${s.engine}`; + if (s.engine.startsWith('ready') || s.engine.startsWith('error')) return; + } catch (e) { /* server booting */ } + await sleep(1500); + } +})(); diff --git a/web/index.html b/web/index.html new file mode 100644 index 0000000..3df44cf --- /dev/null +++ b/web/index.html @@ -0,0 +1,86 @@ + + + + +ROTOGOD + + + + +
+

ROTOGOD

+ + + engine: … + +
+ +
+
+ +
+ +
+
+ + + + + + + – / – +
+
+ +
+
+
+
+ +
+ + + + diff --git a/web/style.css b/web/style.css new file mode 100644 index 0000000..5e7acfe --- /dev/null +++ b/web/style.css @@ -0,0 +1,73 @@ +* { box-sizing: border-box; margin: 0; } +:root { + --bg: #101014; --panel: #17171d; --line: #26262e; + --fg: #e8e8ee; --dim: #8b8b98; --accent: #7ad0ff; +} +body { + background: var(--bg); color: var(--fg); + font: 14px/1.45 -apple-system, "SF Pro Text", Helvetica, sans-serif; + height: 100vh; display: flex; flex-direction: column; overflow: hidden; +} +h1 { font-size: 17px; letter-spacing: 2px; color: var(--accent); } +h1 span { color: var(--fg); } +h2 { font-size: 11px; text-transform: uppercase; letter-spacing: 1.5px; color: var(--dim); margin-bottom: 8px; } +.dim { color: var(--dim); font-size: 12px; } +.hint { font-size: 11px; } +kbd { background: var(--line); border-radius: 3px; padding: 0 4px; font-size: 10px; } + +header { + display: flex; align-items: center; gap: 10px; + padding: 8px 14px; border-bottom: 1px solid var(--line); background: var(--panel); +} +header input#path { flex: 1; } +input, button { + background: #1f1f27; color: var(--fg); border: 1px solid var(--line); + border-radius: 6px; padding: 6px 10px; font-size: 13px; +} +input:focus { outline: 1px solid var(--accent); } +button { cursor: pointer; white-space: nowrap; } +button:hover { border-color: var(--accent); } +button.primary, #startRoto, #track { background: #14303f; border-color: #2b5a75; } +.badge { font-size: 11px; color: var(--dim); border: 1px solid var(--line); border-radius: 10px; padding: 2px 8px; } +#toast { font-size: 12px; color: var(--accent); transition: opacity .4s; } +#toast.err { color: #ff8a80; } + +main { flex: 1; display: flex; min-height: 0; } +#stage { + flex: 1; display: flex; align-items: center; justify-content: center; + background: #000; min-width: 0; +} +#cv { max-width: 100%; max-height: 100%; cursor: crosshair; } + +aside { + width: 265px; overflow-y: auto; border-left: 1px solid var(--line); + background: var(--panel); padding: 12px; +} +aside section { margin-bottom: 18px; } +.row { display: flex; gap: 6px; align-items: center; margin: 6px 0; } +.row.wrap { flex-wrap: wrap; } +.wide { flex: 1; } +progress { width: 100%; height: 6px; accent-color: var(--accent); } +#feather { width: 56px; } + +.obj { + display: flex; align-items: center; gap: 6px; padding: 4px 9px; + border: 1px solid var(--line); border-radius: 14px; cursor: pointer; font-size: 12px; +} +.obj.active { border-color: var(--accent); background: #14303f; } +.obj .swatch { width: 10px; height: 10px; border-radius: 50%; } + +#outlist div { padding: 3px 0; border-bottom: 1px dotted var(--line); word-break: break-all; font-size: 11px; } + +footer { border-top: 1px solid var(--line); background: var(--panel); padding: 6px 14px 10px; } +#transport { display: flex; align-items: center; gap: 6px; margin-bottom: 6px; } +#transport .gap { width: 18px; } +#frameLabel { margin-left: auto; font-variant-numeric: tabular-nums; } + +#timeline { position: relative; height: 44px; margin-bottom: 4px; user-select: none; } +#thumbs { width: 100%; height: 100%; object-fit: fill; border-radius: 4px; opacity: .9; background: #000; } +.shade { position: absolute; top: 0; height: 100%; background: rgba(10,10,14,.72); pointer-events: none; } +#shadeL { left: 0; border-right: 2px solid var(--accent); } +#shadeR { right: 0; border-left: 2px solid var(--accent); } +#playhead { position: absolute; top: -2px; width: 2px; height: calc(100% + 4px); background: #fff; pointer-events: none; } +#scrub { width: 100%; padding: 0; accent-color: var(--accent); }