feat(agent): desktop GPU agent container — CCIP + figure crops over HTTP (#114)
The last piece: a Dockerised desktop-GPU worker that talks to FC ONLY over HTTP (lease → fetch pixels → detect figures + CCIP-embed → submit), so Redis/Postgres stay private. New top-level agent/ (outside CI scope — verified by running it): - fc_agent/worker.py: the lease/compute/submit loop, concurrency 1, start/pause/ stop (stop frees the card; unprocessed leases expire + re-queue). - fc_agent/models.py: imgutils wrappers — detect_person (figures) + CCIP embed. The two API seams to verify against the installed dghs-imgutils (flagged). - fc_agent/media.py: stills + video frame sampling (ffmpeg) at FC's cadence → per-frame instances (the bag). - fc_agent/crops.py: vendored crop primitive. client.py: the FC HTTP client. - fc_agent/app.py: FastAPI localhost control UI (start/pause/stop + progress + queue depth). Dockerfile (CUDA + onnxruntime-gpu + ffmpeg) + requirements + README (token → build → run --gpus all → Start; CPU-fallback path). This completes the CCIP pipeline end to end: agent produces region CCIP vectors → RegionService stores → matcher suggests characters → rail. Verified by running on the desktop (not CI). README calls out the imgutils API + model-string checks. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Ttrj5P7upUTueSfoJcxEqa
This commit is contained in:
@@ -0,0 +1,20 @@
|
|||||||
|
# FabledCurator GPU agent — runs on the desktop with the GPU.
|
||||||
|
# CUDA runtime so onnxruntime-gpu can use the card; ffmpeg for video frames.
|
||||||
|
FROM nvidia/cuda:12.4.1-runtime-ubuntu22.04
|
||||||
|
|
||||||
|
ENV DEBIAN_FRONTEND=noninteractive PYTHONUNBUFFERED=1
|
||||||
|
RUN apt-get update \
|
||||||
|
&& apt-get install -y --no-install-recommends python3 python3-pip ffmpeg \
|
||||||
|
&& rm -rf /var/lib/apt/lists/*
|
||||||
|
|
||||||
|
WORKDIR /app
|
||||||
|
COPY requirements.txt .
|
||||||
|
RUN pip3 install --no-cache-dir -r requirements.txt
|
||||||
|
COPY fc_agent ./fc_agent
|
||||||
|
|
||||||
|
# imgutils caches downloaded ONNX models here; mount a volume to persist them.
|
||||||
|
ENV HF_HOME=/models
|
||||||
|
EXPOSE 8770
|
||||||
|
|
||||||
|
# The control UI; the worker is started from it (or POST /start).
|
||||||
|
CMD ["uvicorn", "fc_agent.app:app", "--host", "0.0.0.0", "--port", "8770"]
|
||||||
@@ -0,0 +1,60 @@
|
|||||||
|
# FabledCurator GPU agent
|
||||||
|
|
||||||
|
A desktop-GPU worker that embeds characters (CCIP) + figure crops for
|
||||||
|
FabledCurator. It talks to FC **only over HTTP** — it leases jobs, fetches image
|
||||||
|
pixels, runs the models on your GPU, and posts results back. Your FC database and
|
||||||
|
Redis stay private; the agent never touches them.
|
||||||
|
|
||||||
|
You run it when you want a burst and stop it to reclaim the card.
|
||||||
|
|
||||||
|
## 1. Get a token
|
||||||
|
In FC: **Settings → Tagging → GPU agent → Generate token** (or Rotate). Copy it.
|
||||||
|
|
||||||
|
## 2. Build
|
||||||
|
```sh
|
||||||
|
cd agent
|
||||||
|
docker build -t fc-gpu-agent .
|
||||||
|
```
|
||||||
|
|
||||||
|
## 3. Run (on the machine with the GPU)
|
||||||
|
```sh
|
||||||
|
docker run --rm --gpus all -p 8770:8770 \
|
||||||
|
-e FC_URL=http://curator.traefik.internal \
|
||||||
|
-e FC_TOKEN=<paste-the-token> \
|
||||||
|
-v fc-agent-models:/models \
|
||||||
|
fc-gpu-agent
|
||||||
|
```
|
||||||
|
Then open <http://localhost:8770> — the control page. Click **Start** to begin
|
||||||
|
draining the queue; **Pause**/**Stop** to yield the GPU. The `-v fc-agent-models`
|
||||||
|
volume caches the downloaded ONNX models so restarts are fast.
|
||||||
|
|
||||||
|
Kick off a backfill from FC (**GPU agent card → Queue character embedding**), then
|
||||||
|
watch the queue counts on the control page (or FC's card) drain.
|
||||||
|
|
||||||
|
## Config (env)
|
||||||
|
| var | default | meaning |
|
||||||
|
|---|---|---|
|
||||||
|
| `FC_URL` | `http://localhost:8000` | FC base URL |
|
||||||
|
| `FC_TOKEN` | — | the bearer token (required) |
|
||||||
|
| `AGENT_ID` | `desktop-agent` | identifies this agent's leases |
|
||||||
|
| `BATCH_SIZE` | `4` | jobs leased per round (still processed one at a time) |
|
||||||
|
| `CCIP_MODEL` | imgutils default | CCIP model name |
|
||||||
|
| `DETECTOR_LEVEL` | `m` | person-detector size: `n` < `s` < `m` < `x` |
|
||||||
|
| `POLL_IDLE_SECONDS` | `10` | wait between empty leases |
|
||||||
|
|
||||||
|
## ⚠️ Verify on first run
|
||||||
|
This part can't be CI-tested (no GPU/models in CI), so confirm against your
|
||||||
|
installed `dghs-imgutils` (`pip show dghs-imgutils`) — see `fc_agent/models.py`:
|
||||||
|
- `imgutils.detect.detect_person(image, level=...)` returns
|
||||||
|
`[((x0,y0,x1,y1), label, score), ...]`.
|
||||||
|
- `imgutils.metrics.ccip_extract_feature(image, model=...)` returns a vector
|
||||||
|
(768-d for caformer). If you want the F1-0.94 variant, set
|
||||||
|
`CCIP_MODEL=ccip-caformer_b36-24` (verify the exact string in imgutils).
|
||||||
|
|
||||||
|
If FC's matcher under/over-fires, tune the cosine threshold in
|
||||||
|
`backend/app/services/ml/ccip.py` (`DEFAULT_SIM_THRESHOLD`) and use
|
||||||
|
`GET /api/ccip/overview` + `/api/ccip/images/<id>` to spot-check.
|
||||||
|
|
||||||
|
## CPU fallback
|
||||||
|
Swap `onnxruntime-gpu` → `onnxruntime` in `requirements.txt` and drop `--gpus all`
|
||||||
|
to grind it slowly on the server instead. Same agent, no card.
|
||||||
@@ -0,0 +1,94 @@
|
|||||||
|
"""FastAPI control surface for the agent (served on localhost).
|
||||||
|
|
||||||
|
Start / pause / resume / stop the worker, set nothing else here (config is env),
|
||||||
|
and watch progress + the server-side queue. The container exposes this on a
|
||||||
|
localhost port; stopping the worker frees the GPU.
|
||||||
|
"""
|
||||||
|
from fastapi import FastAPI
|
||||||
|
from fastapi.responses import HTMLResponse, JSONResponse
|
||||||
|
|
||||||
|
from .config import Config
|
||||||
|
from .worker import Worker
|
||||||
|
|
||||||
|
cfg = Config.from_env()
|
||||||
|
worker = Worker(cfg)
|
||||||
|
app = FastAPI(title="FabledCurator GPU agent")
|
||||||
|
|
||||||
|
|
||||||
|
@app.get("/", response_class=HTMLResponse)
|
||||||
|
def index() -> str:
|
||||||
|
return _PAGE
|
||||||
|
|
||||||
|
|
||||||
|
@app.post("/start")
|
||||||
|
def start():
|
||||||
|
worker.start()
|
||||||
|
return JSONResponse(worker.status())
|
||||||
|
|
||||||
|
|
||||||
|
@app.post("/pause")
|
||||||
|
def pause():
|
||||||
|
worker.pause()
|
||||||
|
return JSONResponse(worker.status())
|
||||||
|
|
||||||
|
|
||||||
|
@app.post("/resume")
|
||||||
|
def resume():
|
||||||
|
worker.resume()
|
||||||
|
return JSONResponse(worker.status())
|
||||||
|
|
||||||
|
|
||||||
|
@app.post("/stop")
|
||||||
|
def stop():
|
||||||
|
worker.stop()
|
||||||
|
return JSONResponse(worker.status())
|
||||||
|
|
||||||
|
|
||||||
|
@app.get("/status")
|
||||||
|
def status():
|
||||||
|
s = worker.status()
|
||||||
|
s["fc_url"] = cfg.fc_url
|
||||||
|
s["configured"] = bool(cfg.token)
|
||||||
|
try:
|
||||||
|
s["queue"] = worker.client.queue_status()
|
||||||
|
except Exception:
|
||||||
|
s["queue"] = None
|
||||||
|
return JSONResponse(s)
|
||||||
|
|
||||||
|
|
||||||
|
_PAGE = """<!doctype html><html><head><meta charset=utf-8>
|
||||||
|
<title>FabledCurator GPU agent</title>
|
||||||
|
<style>
|
||||||
|
body{font:14px system-ui;margin:2rem;max-width:640px;background:#14171a;color:#e8e8e8}
|
||||||
|
h1{font-size:18px} button{font:14px system-ui;padding:.5rem 1rem;border:0;border-radius:6px;
|
||||||
|
margin-right:.5rem;cursor:pointer;color:#fff} .start{background:#2e7d32}.pause{background:#b26a00}
|
||||||
|
.stop{background:#b3261e} .stat{display:inline-block;margin-right:1.5rem}
|
||||||
|
.n{font-size:22px;font-weight:700} code{background:#222;padding:2px 6px;border-radius:4px}
|
||||||
|
.q{margin-top:1rem;color:#9aa}
|
||||||
|
</style></head><body>
|
||||||
|
<h1>FabledCurator GPU agent</h1>
|
||||||
|
<p>FC: <code id=fc>—</code> · token <code id=cfg>—</code></p>
|
||||||
|
<p>
|
||||||
|
<button class=start onclick=act('start')>Start</button>
|
||||||
|
<button class=pause onclick=act('pause')>Pause</button>
|
||||||
|
<button class=pause onclick=act('resume')>Resume</button>
|
||||||
|
<button class=stop onclick=act('stop')>Stop</button>
|
||||||
|
</p>
|
||||||
|
<p>
|
||||||
|
<span class=stat><span class=n id=state>idle</span><br>state</span>
|
||||||
|
<span class=stat><span class=n id=done>0</span><br>processed</span>
|
||||||
|
<span class=stat><span class=n id=err>0</span><br>errors</span>
|
||||||
|
<span class=stat><span class=n id=cur>—</span><br>current image</span>
|
||||||
|
</p>
|
||||||
|
<div class=q id=queue></div>
|
||||||
|
<script>
|
||||||
|
async function act(p){await fetch('/'+p,{method:'POST'});refresh()}
|
||||||
|
async function refresh(){
|
||||||
|
const s=await (await fetch('/status')).json()
|
||||||
|
state.textContent=s.state; done.textContent=s.processed; err.textContent=s.errors
|
||||||
|
cur.textContent=s.current??'—'; fc.textContent=s.fc_url
|
||||||
|
cfg.textContent=s.configured?'set':'MISSING'
|
||||||
|
queue.textContent=s.queue?`queue — pending ${s.queue.pending} · in flight ${s.queue.leased} · done ${s.queue.done} · errored ${s.queue.error}`:'queue — unreachable'
|
||||||
|
}
|
||||||
|
refresh(); setInterval(refresh,3000)
|
||||||
|
</script></body></html>"""
|
||||||
@@ -0,0 +1,66 @@
|
|||||||
|
"""HTTP client for the FabledCurator GPU-job API.
|
||||||
|
|
||||||
|
The agent's ONLY contact with FC — lease/submit/heartbeat/fail + fetch image
|
||||||
|
bytes, all over HTTP with the bearer token. No DB/Redis.
|
||||||
|
"""
|
||||||
|
import requests
|
||||||
|
|
||||||
|
|
||||||
|
class FcClient:
|
||||||
|
def __init__(self, base_url: str, token: str, agent_id: str):
|
||||||
|
self.base = base_url.rstrip("/")
|
||||||
|
self.agent_id = agent_id
|
||||||
|
self.s = requests.Session()
|
||||||
|
self.s.headers["Authorization"] = f"Bearer {token}"
|
||||||
|
|
||||||
|
def lease(self, batch_size: int) -> list[dict]:
|
||||||
|
r = self.s.post(
|
||||||
|
f"{self.base}/api/gpu/jobs/lease",
|
||||||
|
json={"agent_id": self.agent_id, "batch_size": batch_size},
|
||||||
|
timeout=30,
|
||||||
|
)
|
||||||
|
r.raise_for_status()
|
||||||
|
return r.json().get("jobs", [])
|
||||||
|
|
||||||
|
def submit(self, job_id: int, regions: list[dict], replace_kinds: list[str]) -> dict:
|
||||||
|
r = self.s.post(
|
||||||
|
f"{self.base}/api/gpu/jobs/submit",
|
||||||
|
json={
|
||||||
|
"agent_id": self.agent_id, "job_id": job_id,
|
||||||
|
"regions": regions, "replace_kinds": replace_kinds,
|
||||||
|
},
|
||||||
|
timeout=120,
|
||||||
|
)
|
||||||
|
r.raise_for_status()
|
||||||
|
return r.json()
|
||||||
|
|
||||||
|
def heartbeat(self, job_ids: list[int]) -> None:
|
||||||
|
try:
|
||||||
|
self.s.post(
|
||||||
|
f"{self.base}/api/gpu/jobs/heartbeat",
|
||||||
|
json={"agent_id": self.agent_id, "job_ids": job_ids},
|
||||||
|
timeout=30,
|
||||||
|
)
|
||||||
|
except requests.RequestException:
|
||||||
|
pass
|
||||||
|
|
||||||
|
def fail(self, job_id: int, error: str) -> None:
|
||||||
|
try:
|
||||||
|
self.s.post(
|
||||||
|
f"{self.base}/api/gpu/jobs/fail",
|
||||||
|
json={"agent_id": self.agent_id, "job_id": job_id, "error": error},
|
||||||
|
timeout=30,
|
||||||
|
)
|
||||||
|
except requests.RequestException:
|
||||||
|
pass
|
||||||
|
|
||||||
|
def fetch_image(self, image_url: str) -> bytes:
|
||||||
|
# image_url is a server-relative path ("/images/...").
|
||||||
|
r = self.s.get(f"{self.base}{image_url}", timeout=180)
|
||||||
|
r.raise_for_status()
|
||||||
|
return r.content
|
||||||
|
|
||||||
|
def queue_status(self) -> dict:
|
||||||
|
r = self.s.get(f"{self.base}/api/gpu/status", timeout=15)
|
||||||
|
r.raise_for_status()
|
||||||
|
return r.json()
|
||||||
@@ -0,0 +1,26 @@
|
|||||||
|
"""Agent config, all from env (the control container is configured at run)."""
|
||||||
|
import os
|
||||||
|
from dataclasses import dataclass
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass
|
||||||
|
class Config:
|
||||||
|
fc_url: str # base URL of the FabledCurator web service
|
||||||
|
token: str # the bearer token from Settings → Tagging → GPU agent
|
||||||
|
agent_id: str # identifies this agent's leases
|
||||||
|
batch_size: int # jobs leased per round (concurrency is still 1)
|
||||||
|
ccip_model: str # imgutils CCIP model name ("" → imgutils default)
|
||||||
|
detector_level: str # imgutils person-detector level: n|s|m|x
|
||||||
|
poll_idle_seconds: float # wait between empty leases
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
def from_env(cls) -> "Config":
|
||||||
|
return cls(
|
||||||
|
fc_url=os.environ.get("FC_URL", "http://localhost:8000").rstrip("/"),
|
||||||
|
token=os.environ.get("FC_TOKEN", ""),
|
||||||
|
agent_id=os.environ.get("AGENT_ID", "desktop-agent"),
|
||||||
|
batch_size=int(os.environ.get("BATCH_SIZE", "4")),
|
||||||
|
ccip_model=os.environ.get("CCIP_MODEL", ""),
|
||||||
|
detector_level=os.environ.get("DETECTOR_LEVEL", "m"),
|
||||||
|
poll_idle_seconds=float(os.environ.get("POLL_IDLE_SECONDS", "10")),
|
||||||
|
)
|
||||||
@@ -0,0 +1,36 @@
|
|||||||
|
"""Crop primitive — vendored from backend/app/services/ml/crops.py so the agent
|
||||||
|
is self-contained. Keep in sync if the floor logic changes."""
|
||||||
|
from PIL import Image
|
||||||
|
|
||||||
|
MIN_CROP_FRACTION = 0.10
|
||||||
|
MIN_CROP_PX = 64
|
||||||
|
|
||||||
|
|
||||||
|
def crop_region(
|
||||||
|
img: Image.Image,
|
||||||
|
bbox: tuple[float, float, float, float],
|
||||||
|
*,
|
||||||
|
pad: float = 0.0,
|
||||||
|
min_fraction: float = MIN_CROP_FRACTION,
|
||||||
|
min_px: int = MIN_CROP_PX,
|
||||||
|
) -> Image.Image | None:
|
||||||
|
"""Crop a NORMALIZED bbox (x, y, w, h in [0,1]); None if below the size
|
||||||
|
floor (max of a fraction-of-short-side and an absolute pixel floor)."""
|
||||||
|
iw, ih = img.size
|
||||||
|
x, y, w, h = bbox
|
||||||
|
px, py, pw, ph = x * iw, y * ih, w * iw, h * ih
|
||||||
|
if pad:
|
||||||
|
px -= pw * pad / 2.0
|
||||||
|
py -= ph * pad / 2.0
|
||||||
|
pw *= (1.0 + pad)
|
||||||
|
ph *= (1.0 + pad)
|
||||||
|
left = max(0, int(round(px)))
|
||||||
|
top = max(0, int(round(py)))
|
||||||
|
right = min(iw, int(round(px + pw)))
|
||||||
|
bottom = min(ih, int(round(py + ph)))
|
||||||
|
if right <= left or bottom <= top:
|
||||||
|
return None
|
||||||
|
floor = max(min_px, int(min_fraction * min(iw, ih)))
|
||||||
|
if min(right - left, bottom - top) < floor:
|
||||||
|
return None
|
||||||
|
return img.crop((left, top, right, bottom)).convert("RGB")
|
||||||
@@ -0,0 +1,48 @@
|
|||||||
|
"""Image + video handling. Stills load directly; videos are sampled into frames
|
||||||
|
(ffmpeg) at the cadence FC sends — so a video becomes a bag of per-frame
|
||||||
|
instances, each with a timestamp."""
|
||||||
|
import io
|
||||||
|
import os
|
||||||
|
import subprocess
|
||||||
|
import tempfile
|
||||||
|
|
||||||
|
from PIL import Image
|
||||||
|
|
||||||
|
|
||||||
|
def is_video(mime: str) -> bool:
|
||||||
|
return bool(mime) and (mime.startswith("video/") or mime in {"image/gif"})
|
||||||
|
|
||||||
|
|
||||||
|
def load_image(data: bytes) -> Image.Image:
|
||||||
|
return Image.open(io.BytesIO(data)).convert("RGB")
|
||||||
|
|
||||||
|
|
||||||
|
def sample_frames(
|
||||||
|
data: bytes, interval_seconds: float, max_frames: int
|
||||||
|
) -> list[tuple[float, Image.Image]]:
|
||||||
|
"""Extract up to max_frames frames at one-every-interval_seconds via ffmpeg.
|
||||||
|
Returns [(timestamp_seconds, frame)]. Empty on failure (caller falls back)."""
|
||||||
|
interval = max(0.5, float(interval_seconds or 4.0))
|
||||||
|
cap = max(1, int(max_frames or 64))
|
||||||
|
with tempfile.TemporaryDirectory() as tmp:
|
||||||
|
src = os.path.join(tmp, "in")
|
||||||
|
with open(src, "wb") as fh:
|
||||||
|
fh.write(data)
|
||||||
|
pattern = os.path.join(tmp, "f_%05d.jpg")
|
||||||
|
try:
|
||||||
|
subprocess.run(
|
||||||
|
[
|
||||||
|
"ffmpeg", "-nostdin", "-loglevel", "error", "-i", src,
|
||||||
|
"-vf", f"fps=1/{interval}", "-frames:v", str(cap),
|
||||||
|
"-q:v", "3", pattern,
|
||||||
|
],
|
||||||
|
check=True, timeout=600,
|
||||||
|
)
|
||||||
|
except (subprocess.SubprocessError, FileNotFoundError):
|
||||||
|
return []
|
||||||
|
out: list[tuple[float, Image.Image]] = []
|
||||||
|
names = sorted(n for n in os.listdir(tmp) if n.startswith("f_"))
|
||||||
|
for i, name in enumerate(names[:cap]):
|
||||||
|
with Image.open(os.path.join(tmp, name)) as im:
|
||||||
|
out.append((round(i * interval, 2), im.convert("RGB")))
|
||||||
|
return out
|
||||||
@@ -0,0 +1,39 @@
|
|||||||
|
"""imgutils model wrappers — the figure DETECTOR + the CCIP EMBEDDER.
|
||||||
|
|
||||||
|
⚠️ VERIFY ON FIRST RUN: the exact imgutils function names/signatures + the CCIP
|
||||||
|
model string can drift between dghs-imgutils releases. These are the two seams to
|
||||||
|
check against your installed version (`pip show dghs-imgutils`):
|
||||||
|
- detect_person(image, level=...) -> [((x0,y0,x1,y1), label, score), ...]
|
||||||
|
- ccip_extract_feature(image, model=...) -> a vector (768-d for caformer)
|
||||||
|
imgutils auto-downloads the ONNX models from HuggingFace on first use; GPU is
|
||||||
|
used when onnxruntime-gpu is installed.
|
||||||
|
"""
|
||||||
|
import numpy as np
|
||||||
|
from PIL import Image
|
||||||
|
|
||||||
|
|
||||||
|
def detect_figures(image: Image.Image, level: str = "m") -> list[tuple[tuple, float | None]]:
|
||||||
|
"""Person/figure bounding boxes, NORMALIZED (x, y, w, h in [0,1]) + score.
|
||||||
|
Returns [] if detection finds nothing (caller falls back to whole-image)."""
|
||||||
|
from imgutils.detect import detect_person
|
||||||
|
|
||||||
|
iw, ih = image.size
|
||||||
|
out = []
|
||||||
|
for (x0, y0, x1, y1), _label, score in detect_person(image, level=level):
|
||||||
|
out.append((
|
||||||
|
(x0 / iw, y0 / ih, (x1 - x0) / iw, (y1 - y0) / ih),
|
||||||
|
float(score),
|
||||||
|
))
|
||||||
|
return out
|
||||||
|
|
||||||
|
|
||||||
|
def ccip_vector(image: Image.Image, model: str | None = None) -> list[float]:
|
||||||
|
"""The CCIP identity embedding of a (cropped) character image, as a plain
|
||||||
|
float list ready to POST."""
|
||||||
|
from imgutils.metrics import ccip_extract_feature
|
||||||
|
|
||||||
|
feat = (
|
||||||
|
ccip_extract_feature(image, model=model)
|
||||||
|
if model else ccip_extract_feature(image)
|
||||||
|
)
|
||||||
|
return np.asarray(feat, dtype=np.float32).reshape(-1).tolist()
|
||||||
@@ -0,0 +1,127 @@
|
|||||||
|
"""The lease → fetch → detect+embed → submit loop, with start/pause/stop control.
|
||||||
|
|
||||||
|
Concurrency is 1 (one image at a time) so the GPU footprint stays small and a
|
||||||
|
stop frees the card promptly. Stop halts leasing + finishes the current item;
|
||||||
|
unprocessed leases expire and the server re-queues them — nothing is lost.
|
||||||
|
"""
|
||||||
|
import threading
|
||||||
|
import time
|
||||||
|
|
||||||
|
from . import media, models
|
||||||
|
from .client import FcClient
|
||||||
|
from .config import Config
|
||||||
|
from .crops import crop_region
|
||||||
|
|
||||||
|
|
||||||
|
class Worker:
|
||||||
|
def __init__(self, cfg: Config):
|
||||||
|
self.cfg = cfg
|
||||||
|
self.client = FcClient(cfg.fc_url, cfg.token, cfg.agent_id)
|
||||||
|
self._state = "idle" # idle | running | paused | stopping
|
||||||
|
self._lock = threading.Lock()
|
||||||
|
self._thread: threading.Thread | None = None
|
||||||
|
self.processed = 0
|
||||||
|
self.errors = 0
|
||||||
|
self.current = None
|
||||||
|
|
||||||
|
# --- control -----------------------------------------------------------
|
||||||
|
def start(self):
|
||||||
|
with self._lock:
|
||||||
|
if self._state in ("running", "paused"):
|
||||||
|
self._state = "running"
|
||||||
|
return
|
||||||
|
self._state = "running"
|
||||||
|
self._thread = threading.Thread(target=self._run, daemon=True)
|
||||||
|
self._thread.start()
|
||||||
|
|
||||||
|
def pause(self):
|
||||||
|
with self._lock:
|
||||||
|
if self._state == "running":
|
||||||
|
self._state = "paused"
|
||||||
|
|
||||||
|
def resume(self):
|
||||||
|
with self._lock:
|
||||||
|
if self._state == "paused":
|
||||||
|
self._state = "running"
|
||||||
|
|
||||||
|
def stop(self):
|
||||||
|
with self._lock:
|
||||||
|
if self._state in ("running", "paused"):
|
||||||
|
self._state = "stopping"
|
||||||
|
|
||||||
|
def status(self) -> dict:
|
||||||
|
with self._lock:
|
||||||
|
state = self._state
|
||||||
|
return {
|
||||||
|
"state": state, "processed": self.processed,
|
||||||
|
"errors": self.errors, "current": self.current,
|
||||||
|
}
|
||||||
|
|
||||||
|
# --- loop --------------------------------------------------------------
|
||||||
|
def _run(self):
|
||||||
|
while True:
|
||||||
|
with self._lock:
|
||||||
|
st = self._state
|
||||||
|
if st == "stopping":
|
||||||
|
break
|
||||||
|
if st == "paused":
|
||||||
|
time.sleep(1)
|
||||||
|
continue
|
||||||
|
try:
|
||||||
|
jobs = self.client.lease(self.cfg.batch_size)
|
||||||
|
except Exception:
|
||||||
|
time.sleep(self.cfg.poll_idle_seconds)
|
||||||
|
continue
|
||||||
|
if not jobs:
|
||||||
|
time.sleep(self.cfg.poll_idle_seconds)
|
||||||
|
continue
|
||||||
|
ids = [j["job_id"] for j in jobs]
|
||||||
|
for job in jobs:
|
||||||
|
with self._lock:
|
||||||
|
if self._state == "stopping":
|
||||||
|
break
|
||||||
|
self._process(job)
|
||||||
|
self.client.heartbeat(ids) # keep the rest of the batch alive
|
||||||
|
with self._lock:
|
||||||
|
self._state = "idle"
|
||||||
|
|
||||||
|
def _process(self, job: dict):
|
||||||
|
self.current = job.get("image_id")
|
||||||
|
try:
|
||||||
|
data = self.client.fetch_image(job["image_url"])
|
||||||
|
if media.is_video(job.get("mime", "")):
|
||||||
|
frames = media.sample_frames(
|
||||||
|
data, job.get("frame_interval_seconds", 4.0),
|
||||||
|
job.get("max_frames", 64),
|
||||||
|
) or [(None, media.load_image(data))]
|
||||||
|
else:
|
||||||
|
frames = [(None, media.load_image(data))]
|
||||||
|
|
||||||
|
regions = []
|
||||||
|
ev = self.cfg.ccip_model or "ccip-default"
|
||||||
|
dv = f"person-{self.cfg.detector_level}"
|
||||||
|
for t, frame in frames:
|
||||||
|
figs = models.detect_figures(frame, self.cfg.detector_level)
|
||||||
|
if not figs:
|
||||||
|
figs = [((0.0, 0.0, 1.0, 1.0), None)] # whole-frame fallback
|
||||||
|
for bbox, score in figs:
|
||||||
|
crop = crop_region(frame, bbox)
|
||||||
|
if crop is None:
|
||||||
|
continue
|
||||||
|
vec = models.ccip_vector(crop, self.cfg.ccip_model or None)
|
||||||
|
regions.append({
|
||||||
|
"kind": "figure",
|
||||||
|
"bbox": list(bbox),
|
||||||
|
"frame_time": t,
|
||||||
|
"score": score,
|
||||||
|
"ccip_embedding": vec,
|
||||||
|
"embedding_version": ev,
|
||||||
|
"detector_version": dv,
|
||||||
|
})
|
||||||
|
self.client.submit(job["job_id"], regions, ["figure", "face"])
|
||||||
|
self.processed += 1
|
||||||
|
except Exception as exc: # noqa: BLE001 — report + move on
|
||||||
|
self.errors += 1
|
||||||
|
self.client.fail(job["job_id"], str(exc)[:500])
|
||||||
|
finally:
|
||||||
|
self.current = None
|
||||||
@@ -0,0 +1,11 @@
|
|||||||
|
# CCIP + figure detection (ONNX models, auto-downloaded from HuggingFace).
|
||||||
|
dghs-imgutils>=0.4
|
||||||
|
# GPU inference for the ONNX models. Swap to onnxruntime (CPU) for a slow
|
||||||
|
# server-side fallback run.
|
||||||
|
onnxruntime-gpu
|
||||||
|
# Control surface + HTTP.
|
||||||
|
fastapi
|
||||||
|
uvicorn[standard]
|
||||||
|
requests
|
||||||
|
pillow
|
||||||
|
numpy
|
||||||
Reference in New Issue
Block a user