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 <noreply@anthropic.com>
This commit is contained in:
m3ultra 2026-07-26 22:02:20 +10:00
commit 874aaaefa3
15 changed files with 1285 additions and 0 deletions

6
.gitignore vendored Normal file
View File

@ -0,0 +1,6 @@
.venv/
__pycache__/
*.pyc
output/
samples/
.DS_Store

68
DESIGN.md Normal file
View File

@ -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). 84748479 are already taken on this box. Browser canvas is plenty fast for 7201080p 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.

22
README.md Normal file
View File

@ -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).

View File

@ -0,0 +1,27 @@
<?xml version="1.0" encoding="UTF-8"?>
<!--
Always-on rotogod on m3ultra, bound to the TAILNET IP only (LAN/public never
see it; digalot.fyi/rotogod/ reaches it through the dealgod-nginx proxy).
Install:
cp deploy/rotogod.launchd.plist ~/Library/LaunchAgents/party.monster.rotogod.plist
launchctl bootstrap gui/$(id -u) ~/Library/LaunchAgents/party.monster.rotogod.plist
Logs: ~/Library/Logs/rotogod.log
-->
<!DOCTYPE plist PUBLIC "-//Apple//DTD PLIST 1.0//EN" "http://www.apple.com/DTDs/PropertyList-1.0.dtd">
<plist version="1.0">
<dict>
<key>Label</key><string>party.monster.rotogod</string>
<key>ProgramArguments</key>
<array>
<string>/Users/m3ultra/Documents/rotogod/.venv/bin/python</string>
<string>/Users/m3ultra/Documents/rotogod/run.py</string>
<string>--host</string><string>100.89.131.57</string>
<string>--port</string><string>8484</string>
</array>
<key>WorkingDirectory</key><string>/Users/m3ultra/Documents/rotogod</string>
<key>RunAtLoad</key><true/>
<key>KeepAlive</key><true/>
<key>StandardOutPath</key><string>/Users/m3ultra/Library/Logs/rotogod.log</string>
<key>StandardErrorPath</key><string>/Users/m3ultra/Library/Logs/rotogod.log</string>
</dict>
</plist>

10
requirements.txt Normal file
View File

@ -0,0 +1,10 @@
torch
torchvision
transformers
huggingface_hub
fastapi
uvicorn[standard]
av
opencv-python-headless
pillow
numpy

0
rotogod/__init__.py Normal file
View File

88
rotogod/export.py Normal file
View File

@ -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)

44
rotogod/jobs.py Normal file
View File

@ -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

144
rotogod/segment.py Normal file
View File

@ -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

279
rotogod/server.py Normal file
View File

@ -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()}

115
rotogod/video.py Normal file
View File

@ -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

16
run.py Normal file
View File

@ -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")

307
web/app.js Normal file
View File

@ -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 = `<span class="swatch" style="background:${PALETTE[(id - 1) % PALETTE.length]}"></span>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);
}
})();

86
web/index.html Normal file
View File

@ -0,0 +1,86 @@
<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="utf-8">
<title>ROTOGOD</title>
<meta name="viewport" content="width=device-width, initial-scale=1">
<link rel="stylesheet" href="static/style.css">
</head>
<body>
<header>
<h1>ROTO<span>GOD</span></h1>
<input id="path" placeholder="/absolute/path/to/clip.mp4" spellcheck="false">
<button id="load">Load</button>
<span id="engine" class="badge">engine: …</span>
<span id="toast"></span>
</header>
<main>
<div id="stage"><canvas id="cv" width="16" height="9"></canvas></div>
<aside>
<section>
<h2>Roto</h2>
<div id="rangeInfo" class="dim">range: </div>
<div class="row">
<button id="startRoto" class="wide">Start roto on in→out</button>
<button id="resetRoto" title="drop session"></button>
</div>
<div id="objects" class="row wrap"></div>
<div class="row">
<button id="addObj">+ object</button>
<span class="dim hint">L-click add · R/⌥-click subtract</span>
</div>
<div class="row">
<button id="track" class="wide">▶ Track range</button>
<button id="trackHere" title="re-propagate from current frame">▶ here</button>
</div>
<progress id="prog" value="0" max="1"></progress>
<div id="progText" class="dim"></div>
<label class="dim"><input type="checkbox" id="showOverlay" checked> show masks</label>
</section>
<section>
<h2>Export roto</h2>
<div class="row wrap">
<button class="exp" data-kind="prores">ProRes 4444 α</button>
<button class="exp" data-kind="webm">WebM α</button>
<button class="exp" data-kind="matte">Matte</button>
<button class="exp" data-kind="green">Greenscreen</button>
</div>
<label class="dim">feather <input id="feather" type="number" value="2" min="0" max="20"></label>
</section>
<section>
<h2>Scene cut</h2>
<div class="row">
<button id="trim" class="wide">Export in→out</button>
<label class="dim"><input type="checkbox" id="exactTrim" checked> exact</label>
</div>
</section>
<section>
<h2>Output</h2>
<div id="outlist" class="dim"></div>
</section>
</aside>
</main>
<footer>
<div id="transport">
<button id="prev10">«</button><button id="prev"></button>
<button id="play"></button>
<button id="next"></button><button id="next10">»</button>
<span class="gap"></span>
<button id="setIn">⇤ in <kbd>i</kbd></button>
<button id="setOut">out ⇥ <kbd>o</kbd></button>
<span id="frameLabel" class="dim"> / </span>
</div>
<div id="timeline">
<img id="thumbs" draggable="false" alt="">
<div id="shadeL" class="shade"></div>
<div id="shadeR" class="shade"></div>
<div id="playhead"></div>
</div>
<input id="scrub" type="range" min="0" max="0" value="0" step="1">
</footer>
<script src="static/app.js"></script>
</body>
</html>

73
web/style.css Normal file
View File

@ -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); }