96bba8e538
Makes the digalot.fyi/rotogod/ proxy actually usable away from the box:
⬆ uploads land in samples/uploads and open; output list links through
GET /api/download/{name}. nginx side: 8g body cap, request buffering off.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
316 lines
8.9 KiB
Python
316 lines
8.9 KiB
Python
"""rotogod server — FastAPI over one clip + one roto session (single-user tool)."""
|
|
|
|
import contextlib
|
|
import os
|
|
import re
|
|
import threading
|
|
|
|
import cv2
|
|
import numpy as np
|
|
from fastapi import FastAPI, HTTPException, Request
|
|
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()
|
|
|
|
|
|
UPLOAD_CAP = 8 * 2 ** 30 # 8GB — plenty for ProRes sources on this box
|
|
|
|
|
|
@app.post("/api/upload")
|
|
async def upload(request: Request, name: str = "clip.mp4"):
|
|
"""Raw-body clip upload from the browser; lands in samples/uploads and opens."""
|
|
safe = re.sub(r"[^A-Za-z0-9._-]", "_", os.path.basename(name)) or "clip.mp4"
|
|
updir = os.path.join(ROOT, "samples", "uploads")
|
|
os.makedirs(updir, exist_ok=True)
|
|
dest = os.path.join(updir, safe)
|
|
size = 0
|
|
try:
|
|
with open(dest, "wb") as f:
|
|
async for chunk in request.stream():
|
|
size += len(chunk)
|
|
if size > UPLOAD_CAP:
|
|
raise HTTPException(413, "upload too large")
|
|
f.write(chunk)
|
|
except HTTPException:
|
|
os.remove(dest)
|
|
raise
|
|
if size == 0:
|
|
os.remove(dest)
|
|
raise HTTPException(400, "empty upload")
|
|
return open_clip(OpenBody(path=dest))
|
|
|
|
|
|
@app.get("/api/download/{name}")
|
|
def download(name: str):
|
|
p = os.path.join(export.OUTPUT_DIR, os.path.basename(name))
|
|
if not os.path.isfile(p):
|
|
raise HTTPException(404)
|
|
return FileResponse(p, filename=os.path.basename(p))
|
|
|
|
|
|
@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()}
|