diff --git a/.claude/skills/setup-video-segmentation-cvat/SKILL.md b/.claude/skills/setup-video-segmentation-cvat/SKILL.md new file mode 100644 index 00000000..0fc503f6 --- /dev/null +++ b/.claude/skills/setup-video-segmentation-cvat/SKILL.md @@ -0,0 +1,165 @@ +--- +name: setup-video-segmentation-cvat +description: > + Use when working in examples/video_segmentation_cvat, or when setting up / running / debugging the + datapipe video → SAM3 → CVAT example (long egocentric video → sampled frames → text-prompt + segmentation → CVAT pre-annotations), on the built-in city-walk videos or your own footage. +--- + +# video_segmentation_cvat (video → SAM3 → CVAT) + +This skill = turn a long video into review-ready CVAT pre-annotations. Front stage samples + dedups +frames; the SAM3→CVAT tail is identical to [`../sam_cvat`](../sam_cvat). Set the knobs below first. + +**Ask first — don't assume (only the unresolved):** **is this a live demo on the built-in test +videos, or a real run on the user's own data?** (demo → the two-stage choreography below); which videos +(built-in city-walk set via `scripts/fetch_video.py`, or the user's own) → `INPUT_VIDEO_DIR`? +**`SAMPLE_FPS`** (how densely to sample — the single biggest lever on frame count / run time) and +**`SEGMENT_SIZE`** (frames per CVAT job)? **which Postgres + which database** for `DB_URL` — never +point at an existing DB or use a default without confirming; external CVAT ready or provision it? +reuse an existing venv / `uv` env or create fresh? which GPU (VRAM + FlashAttention)? surface stage +logs or run quiet? + +**How to work:** read the setup, propose a short plan, get a go-ahead before touching anything. +Prepare `.env` and **pause for the user to verify it** before running. Run stages with logs shown; +after each, say what you did and what changed — don't run silently. If a stage fails and the cause +isn't clear, re-run with `datapipe --debug … run` (`--debug-sql` for SQL); debug is very verbose, so +send it to a file and `grep` (`datapipe --debug run > /tmp/dp.log 2>&1; grep -nEi "error|traceback" /tmp/dp.log`). + +## Pipeline +``` +stage=video list_videos folder INPUT_VIDEO_DIR -> video +stage=sample extract_frames ffmpeg fps=SAMPLE_FPS -> frames +stage=sample dedup_frames perceptual-hash dedup -> local_images +stage=prompt list_sam_config SAM_TEXT_PROMPT -> sam_config +stage=sam sam_inference SAM3 image-mode -> sam_predictions +stage=sam sam_to_cvat_xml -> sam_cvat_xml +stage=cvat prepare_cvat_input / CVATStep / parse_cvat_annotations -> image__annotations +``` +`local_images (video_id, image_id, image_path)` is what the SAM→CVAT tail consumes; everything from +`sam_inference` on (incl. `models.py`) is a **verbatim copy of `sam_cvat`** — keep it that way. The +video front's only job is video → deduped frames (SAM runs at native resolution). **One CVAT task per +video** (`task_queue_id=video_id`, `FILES_BATCH` huge), split into jobs of `SEGMENT_SIZE` frames. + +## Sampling knobs (the front) +- **`SAMPLE_FPS`** (default 0.2 = 1 frame/5 s). Walking POV is highly redundant — 1 fps vs 1/3 fps + barely differ in what gets annotated; sample sparser to cut annotation load. **Ask before a run.** +- **`PHASH_MAX_DISTANCE`** (Hamming, default 10) — higher = more aggressive near-duplicate dedup. +- **Gotcha:** `extract_frames` reuses frames already in `FRAMES_DIR`, so **changing `SAMPLE_FPS` has + no effect until you clear `FRAMES_DIR`** (`rm -rf /`). +- Align class across all three: `SAM_TEXT_PROMPT` == `CVAT_BOX_LABEL`/`CVAT_POLYGON_LABEL` — mismatch + runs clean but yields 0 useful annotations. Levers: `SAM_SCORE_THRESHOLD` (0.5), `SAM_MAX_DETECTIONS` (20). + +## GPU +SAM3 runs at the frame's **native resolution** (no in-pipeline resize). *Field note (measured):* a +FlashAttention GPU (Ada/Ampere+, 16 GB) does full 1280×720 in **~5.3 GB, ~0.27 s/frame**. A +non-FlashAttention (Pascal-class) 8 GB card OOMs on raw 720p — use a newer/bigger GPU, or shrink +frames upstream (e.g. add a `scale` filter to `extract_frames`) rather than editing the shared tail. + +## Prerequisites +- **GPU** (see above). **HF token** — SAM3 is gated: accept the license at huggingface.co/facebook/sam3, + set `HF_TOKEN` (gated read access). **Read-only home?** set `HF_HOME=/writable/path` so weights + + token cache there (the example already avoids persisting the token to disk when `HF_TOKEN` is set). +- **`ffmpeg`** on `PATH`, else the `imageio-ffmpeg` bundled binary is used automatically. +- **External PostgreSQL** at `DB_URL` (tables auto-create via `datapipe db create-all`). +- **External CVAT** at `CVAT_URL` (user/pass, optional org), project `CVAT_PROJECT_ID` whose labels + match `CVAT_BOX_LABEL` / `CVAT_POLYGON_LABEL` (defaults `person_box` / `person_mask`). +- **`uv` + Python ≥3.10,<3.13** → `uv sync`. Pins cu124 torch, editable local libs + (`../../libs/datapipe-*`, monorepo-only), builds `sam3` from a pinned git rev (+`imagehash`). After + `uv sync`, on a pre-AVX2 host re-apply `uv pip install polars-lts-cpu==1.33.1`. On a very new GPU + whose CUDA arch the pinned torch predates (`CUDA error: no kernel image is available` on a tiny cuda + matmul), reinstall a matching build — for **Blackwell / RTX 50-series use `cu128`** (torch ≥2.7): + `uv pip install --python .venv --reinstall torch torchvision --index-url https://download.pytorch.org/whl/cu128`. + That index can flake (503s) — wrap it in a retry loop. + +## Stand up CVAT + Postgres (if you don't already have them) +Both run in Docker; the example just points `DB_URL` / `CVAT_URL` at them (no code in it manages them). +- **Postgres:** `docker run -d --name dp_pg -e POSTGRES_PASSWORD= -e POSTGRES_DB=postgres -p 5432:5432 postgres:15` → `DB_URL=postgresql+psycopg2://postgres:@localhost:5432/postgres`. +- **CVAT v2.65.0** (matches `cvat-sdk==2.65.0`): `git clone --depth 1 --branch v2.65.0 https://github.com/cvat-ai/cvat.git /opt/cvat && cd /opt/cvat && docker compose up -d` (UI on `:8080`; ~15 containers, needs a few GB RAM + a reliable registry uplink — a flaky link fails the image pulls). The repo's `libs/datapipe-cvat/tests/start-cvat.sh` does the same but `docker compose down`s on exit (test helper), so for a persistent instance run `up -d` directly. +- **Admin** (first run): `docker exec cvat_server bash -lc "DJANGO_SUPERUSER_PASSWORD=admin python3 ~/manage.py createsuperuser --username admin --email a@e.com --noinput"`. +- **Project + labels:** create a project with labels named `CVAT_BOX_LABEL` / `CVAT_POLYGON_LABEL` (default `person_box` / `person_mask`, type `any` accepts box+polygon) and set `CVAT_PROJECT_ID` to its id. Via API: `POST /api/projects` with `{"name":...,"labels":[{"name":"person_box","type":"any"},{"name":"person_mask","type":"any"}]}`. +- Docker-less pod → keep CVAT + Postgres on another host and point `DB_URL`/`CVAT_URL` at it (see the bare-pod section). + +## Get a video +```bash +python scripts/fetch_video.py # whole built-in test set (city walks + smoke clip) +python scripts/fetch_video.py smoke_shibuya_3min # only the smoke clip (stage-2 live demo) +python scripts/fetch_video.py 27Pv4Cg4EV4 ... # specific bucket keys +``` +Pulls pre-encoded 720p `.webm` via `curl` (resumable) from the demo bucket (`$VIDEO_BUCKET_URL`, +default the `e8-demo` bucket). Internal-demo only (source under YouTube ToS) — don't redistribute. +For the two-stage demo: stage 1 fetches the city walks, stage 2 fetches only `smoke_shibuya_3min`. +Own footage → just drop any video file into `INPUT_VIDEO_DIR`. + +## Run +```bash +cp .env.example .env # DB_URL, HF_TOKEN, INPUT_VIDEO_DIR, CVAT_*, SAM_TEXT_PROMPT, SAMPLE_FPS, SEGMENT_SIZE +uv sync && source .venv/bin/activate # else prefix each command with `uv run` +datapipe db create-all && datapipe run +# by stage: datapipe step --labels stage=sample run (then stage=sam, stage=cvat) +``` +Run from `examples/video_segmentation_cvat/` (`app.py` `load_dotenv()`s before importing config). + +## Demo choreography — two stages (built-in test set) +For a live demo, split the built-in videos into a pre-baked bulk and one small live clip, so the +audience sees a run finish in seconds instead of waiting hours: + +- **Stage 1 — bulk, YOU run it fully ahead of time (CLI).** Put every built-in video **except** + `smoke_shibuya_3min` in `INPUT_VIDEO_DIR` and `datapipe run` to completion. This is the slow part + (the multi-hour city-walk videos) and pre-fills CVAT with the heavy tasks. Do it before the demo, + not on stage. Keep the UI **off** during this CLI run (see the `datapipe api` collision below). +- **Stage 2 — smoke clip, LIVE via the UI.** During the demo, start `datapipe api`, drop + `smoke_shibuya_3min` into `INPUT_VIDEO_DIR`, and trigger the run from the **UI's per-stage run + buttons**. It's tiny, so frames → SAM3 → a fresh CVAT task appear within seconds, live — and running + from the UI (single writer) avoids the run-log collision that a concurrent CLI run would cause. + +So: bulk is proof-of-scale (done offline), smoke is the live "watch it work" moment. On the user's own +data there's no split — just one `datapipe run` (or per-video, incrementally). + +## Deploy on a bare GPU pod (no Docker) — field-tested +The example itself needs **no Docker** — it's a plain `uv` venv. Docker is only for the infra deps +(Postgres + CVAT), so on a Docker-less pod either point `DB_URL`/`CVAT_URL` at a host that runs them +(needs network reach to their ports) or, for a SAM-only smoke test, skip them. + +Setup that worked, in order: +1. **Get the code on the branch** without disturbing an existing checkout: from the repo dir, + `git fetch origin ` then `git worktree add ` (reuses the host's git creds, + isolated working tree). **Verify it's current** (`git log origin/..HEAD` empty) — a stale + checkout can ship the old YouTube `fetch_video.py` (needs yt-dlp) instead of the current + bucket-fetch one, and an out-of-date pipeline; `git pull`/re-checkout the file if behind. +2. **`uv sync`** in the example dir. `uv` auto-fetches a Python 3.10–3.12 even if the base interpreter + is older. Takes a few min (torch cu124 ~2.5 GB + builds `sam3`/`cv-pipeliner` from git). **`sam3` + builds with no `nvcc`/CUDA toolkit** — it's pure PyTorch, no compiled CUDA ext (the usual build + worry is a non-issue). +3. **`.env`** — minimum for a SAM run: `HF_TOKEN` (gated), `HF_HOME=/writable/path` (the default home + cache may be small/read-only; weights are ~3.3 GB), `SAM_TEXT_PROMPT`, and **a `DB_URL` even if + unused** — `config.py` builds `DBConn(DB_URL)` at import and crashes on `None` (a dummy + `postgresql+psycopg2://x:x@localhost:5432/x` is enough when you only call `infer_image`). +4. **Run scripts from the example dir** (or set `PYTHONPATH` to it) — `import models`/`config` are + top-level modules, not a package. + +**GPU field note (measured):** on a FlashAttention GPU (Ada/Ampere+, 16 GB), **full 1280×720 SAM3 = +~5.3 GB peak, ~0.27 s/frame** — comfortable, ~10× faster than a non-FA card. First run pays a one-time +~250 s model load + weight download. + +## Debug UI (`datapipe api`) +`datapipe api` (default `:8000`) serves the pipeline graph, table browser, and per-stage run triggers. +**Do not run `datapipe api` and a separate CLI `datapipe run` on the same DB at once** — both write +run logs to the shared Postgres and collide on `uq_run_log_seq` (the API's orphan-run reconciler +adopts the live CLI run), which kills the run. Either trigger runs from the UI, or stop the UI while a +CLI run is in flight. (Radical fix: give the API a dedicated ClickHouse `RunLogsBackend`.) + +## Annotate (human-in-the-loop) +In CVAT fix the pre-annotations, **mark the task completed**, re-run `datapipe run` → edits land in +`image__annotations`. CVAT tables (for wipes): `image_batches`, `cvat_task`, `cvat_images`, +`cvat_task_sync_table`. `CVATStep` does NOT push new pre-annotations to EXISTING tasks (only on image +add/remove/path change) — to change `SAM_TEXT_PROMPT`, wipe old tasks + datapipe CVAT tables, re-run. + +## Troubleshooting (may already be fixed — verify against current files) +- **0 detections** → class misaligned; align `SAM_TEXT_PROMPT` with the CVAT labels. +- **OOM** → non-FlashAttention / too-small GPU for native-res SAM3; use a bigger/newer card or shrink + frames upstream (`scale` filter in `extract_frames`). +- **Changing `SAMPLE_FPS` did nothing** → clear `FRAMES_DIR` (frames are reused). +- **`uq_run_log_seq` duplicate key** → UI + CLI ran together; run one at a time. +- **Read-only filesystem on HF cache** → set `HF_HOME` to a writable path. +- **CVAT rejects label** → project needs labels named `CVAT_BOX_LABEL`/`CVAT_POLYGON_LABEL`. diff --git a/examples/video_segmentation_cvat/.env.example b/examples/video_segmentation_cvat/.env.example new file mode 100644 index 00000000..68772a8e --- /dev/null +++ b/examples/video_segmentation_cvat/.env.example @@ -0,0 +1,38 @@ +DB_URL=postgresql+psycopg2://postgres:postgres@localhost:5432/postgres + +# --- video ingest / frame sampling --- +# Folder you drop long egocentric videos into. Every video here is sampled and segmented. +INPUT_VIDEO_DIR=videos +# Where extracted frames are written (one subfolder per video). Defaults to ./.frames if unset. +# FRAMES_DIR=/tmp/robots-frames +# ffmpeg extraction rate (frames per second) -- ASK how densely to sample before a run. +# Walking POV barely differs between 1 and 1/3 fps; sample sparser to cut annotation load: +# 1 = 1 frame/s, 0.2 = 1 frame/5s, 0.1 = 1 frame/10s. Dedup below removes near-duplicates on top. +SAMPLE_FPS=0.2 +# Perceptual-hash near-duplicate threshold (Hamming). Higher = more aggressive dedup. +PHASH_MAX_DISTANCE=10 + +# --- SAM3 (gated model: accept the license on HuggingFace and set a token) --- +HF_TOKEN=replace-me +SAM_TEXT_PROMPT=person +SAM_SCORE_THRESHOLD=0.5 +# Max detections kept per frame (top-N by score). Crowded street scenes often have >20 people, so a +# low cap drops real detections — 50 suits busy egocentric footage; raise/lower per your scenes. +SAM_MAX_DETECTIONS=50 +# HF caches the gated model + token under $HF_HOME (default ~/.cache/huggingface). Point it at a +# writable path if the home dir is read-only (else login()/weights download fail on such hosts). +# HF_HOME=/var/tmp/hf_home + +# --- CVAT --- +CVAT_URL=http://localhost:8080 +CVAT_USERNAME=admin +CVAT_PASSWORD=admin +CVAT_PROJECT_ID=1 +CVAT_ORGANIZATION= +CVAT_BOX_LABEL=person_box +CVAT_POLYGON_LABEL=person_mask +TASK_QUEUE_ID=queue1 +# One CVAT task per video (task_queue_id=video_id) with jobs of SEGMENT_SIZE frames each. +FILES_BATCH=100000 +MIN_FILES_IN_JOB=1 +SEGMENT_SIZE=200 diff --git a/examples/video_segmentation_cvat/.gitignore b/examples/video_segmentation_cvat/.gitignore new file mode 100644 index 00000000..708821e8 --- /dev/null +++ b/examples/video_segmentation_cvat/.gitignore @@ -0,0 +1,10 @@ +.env +.venv/ +.frames/ +videos/ +*.webm +*.mp4 +*.mkv +*.mov +__pycache__/ +*.pyc diff --git a/examples/video_segmentation_cvat/README.md b/examples/video_segmentation_cvat/README.md new file mode 100644 index 00000000..9d0f21f6 --- /dev/null +++ b/examples/video_segmentation_cvat/README.md @@ -0,0 +1,104 @@ +# Egocentric video → SAM3 → CVAT ("robots" demo) + +Drop a long first-person video into a folder — the pipeline samples frames, segments them with a +**text prompt** (SAM3), and hands you ready-to-review pre-annotations in **CVAT**. Built for a +robotics audience: *"here is the inbox folder, put your 20-hour egocentric recording in it, we +process it for you."* + +It reuses the whole SAM3→CVAT tail from [`../sam_cvat`](../sam_cvat); the only thing added on the +front is turning a long video into a deduplicated set of frames. + +## Pipeline + +``` +stage=video list_videos folder INPUT_VIDEO_DIR -> video +stage=sample extract_frames ffmpeg fps=SAMPLE_FPS -> frames +stage=sample dedup_frames perceptual-hash dedup -> local_images +stage=prompt list_sam_config SAM_TEXT_PROMPT -> sam_config +stage=sam sam_inference SAM3 image-mode -> sam_predictions +stage=sam sam_to_cvat_xml -> sam_cvat_xml +stage=cvat prepare_cvat_input / CVATStep / parse_cvat_annotations -> image__annotations +``` + +`local_images (image_id, image_path)` is exactly what the SAM→CVAT tail consumes, so everything from +`sam_inference` onward is identical to `sam_cvat` — `models.py` is a verbatim copy. The only thing the +video front adds is turning a long video into a deduplicated set of frames. + +## Why sample at 1 fps then dedup (not "every frame", not 1/3 fps) + +Walking POV footage is extremely redundant — you move ~1 m/s and consecutive frames are near +duplicates, so **1 fps and 1/3 fps are not meaningfully different** in what ends up annotated. The +real lever is content-awareness. So: extract at `SAMPLE_FPS=1` for good event coverage, then drop +near-duplicates with a perceptual hash (`PHASH_MAX_DISTANCE`, Hamming). A 24h source at 1 fps is +~86k frames; after dedup it collapses to a few thousand genuinely different frames — enough coverage +without flooding CVAT. Both are ordinary incremental datapipe stages: add another video and only its +frames are processed. + +## Prerequisites + +- **GPU** for SAM3 (`DEVICE` auto-selects `cuda:0`). SAM3 runs at the frame's native resolution — a + FlashAttention-capable card (Ada/Ampere+) handles full 720p in ~5 GB. A non-FlashAttention + (Pascal-class) card falls back to O(n²) math-attention and OOMs 8 GB on 720p; use a bigger/newer GPU + or sample smaller frames. +- **`ffmpeg`** on `PATH` (frame extraction). Not required: the example falls back to the binary + bundled by `imageio-ffmpeg` when no system `ffmpeg` is found. +- **SAM3 is a gated HuggingFace model** — accept the license on the SAM3 model page, create a token, + set `HF_TOKEN` in `.env`. If the home dir is read-only, point `HF_HOME` at a writable path so the + gated weights + token cache there. +- **CVAT** deployed at `CVAT_URL` (see [`../datapipe_cvat/simple_project`](../datapipe_cvat/simple_project/README.md) + for a local Docker setup), a project created (`CVAT_PROJECT_ID`) with labels matching + `CVAT_BOX_LABEL` / `CVAT_POLYGON_LABEL` (defaults `person_box` / `person_mask`). + +## Get a video + +```bash +# whole built-in ~24h test set (busy IN/JP/US city walks + a 3-min smoke clip), into $INPUT_VIDEO_DIR: +python scripts/fetch_video.py +# just the smoke clip (for the live demo): +python scripts/fetch_video.py smoke_shibuya_3min +# specific videos (bucket keys): +python scripts/fetch_video.py 27Pv4Cg4EV4 BsiHD4m6_BU +``` + +`fetch_video.py` pulls the pre-encoded 720p `.webm` files with `curl` (resumable) from the demo +object-storage bucket — override the base URL with `$VIDEO_BUCKET_URL`. For your own footage, drop any +video file into `INPUT_VIDEO_DIR` directly (`.mp4/.mkv/.mov/.webm/.m4v`). + +Internal-demo only — the source footage is under YouTube ToS; do not redistribute the videos or the +frames. For license-clean footage use stock (Pexels/Mixkit) or a dataset clip (EPIC-KITCHENS / Ego4D). + +## Run + +```bash +cp .env.example .env # set DB_URL, HF_TOKEN, INPUT_VIDEO_DIR, CVAT_*, SAM_TEXT_PROMPT +uv sync +datapipe db create-all +datapipe run +``` + +Run a single stage: `datapipe step --labels stage=sample run`, `... stage=sam run`, +`... stage=cvat run`. + +SAM3 runs at the frame's native resolution. **Field note (measured):** on a FlashAttention GPU +(Ada/Ampere+, 16 GB) full 1280×720 SAM3 peaks at only **~5.3 GB, ~0.27 s/frame**. A non-FlashAttention +(Pascal-class) 8 GB card OOMs on 720p — use a newer/bigger GPU, or reduce frame size upstream (e.g. an +`ffmpeg` `scale` filter in `extract_frames`). + +## Debug UI (`datapipe api`) + +`app.py` exposes a `DatapipeAPI`, so `datapipe api` (default `:8000`) serves the pipeline graph, table +browser, and per-stage run triggers. **Do not run `datapipe api` and a separate CLI `datapipe run` +against the same database at the same time** — both write run logs to the shared Postgres and collide +on `uq_run_log_seq` (the API's orphan-run reconciler adopts the live CLI run), which kills the run. +Either trigger runs from the UI, or stop the UI while a CLI run is in flight. + +## Live demo + +Drop a ~10s clip into `INPUT_VIDEO_DIR` and `datapipe run`: it is sampled, deduped, segmented, and a +new CVAT task with box+polygon pre-annotations appears within a few seconds. Change `SAM_TEXT_PROMPT` +(e.g. `person` → `car`) to show open-vocabulary segmentation without retraining. + +## Annotate in CVAT + +Open CVAT, find the task, review/fix the SAM pre-annotations, save, mark completed. Re-run the +pipeline to pull the reviewed annotations back into `image__annotations`. diff --git a/examples/video_segmentation_cvat/__init__.py b/examples/video_segmentation_cvat/__init__.py new file mode 100644 index 00000000..b764ee30 --- /dev/null +++ b/examples/video_segmentation_cvat/__init__.py @@ -0,0 +1 @@ +"""Egocentric video -> frame sampling -> SAM3 -> CVAT annotation pipeline example.""" diff --git a/examples/video_segmentation_cvat/app.py b/examples/video_segmentation_cvat/app.py new file mode 100644 index 00000000..32ff90a9 --- /dev/null +++ b/examples/video_segmentation_cvat/app.py @@ -0,0 +1,121 @@ +from __future__ import annotations + +from dotenv import load_dotenv + +load_dotenv() + +from datapipe.compute import Catalog, Pipeline +from datapipe.datatable import DataStore +from datapipe.executor import ExecutorConfig +from datapipe.step.batch_generate import BatchGenerate +from datapipe.step.batch_transform import BatchTransform +from datapipe_app import DatapipeAPI +from datapipe_cvat.cvat_step import CVATStep + +import data +import steps +from config import ( + CVAT_ORGANIZATION, + CVAT_PASSWORD, + CVAT_PROJECT_ID, + CVAT_URL, + CVAT_USERNAME, + DBCONN, + FILES_BATCH, + MIN_FILES_IN_JOB, + PRIMARY_KEYS, + SEGMENT_SIZE, +) + +pipeline = Pipeline( + [ + # --- video -> frames -------------------------------------------------------------------- + BatchGenerate( + steps.list_videos, + outputs=[data.video_tbl], + labels=[("stage", "video")], + ), + BatchTransform( + func=steps.extract_frames, + inputs=[data.video_tbl], + outputs=[data.frames_tbl], + transform_keys=["video_id"], + chunk_size=1, + labels=[("stage", "sample")], + ), + BatchTransform( + func=steps.dedup_frames, + inputs=[data.frames_tbl], + outputs=[data.local_images_tbl], + transform_keys=["video_id"], + chunk_size=1, + labels=[("stage", "sample")], + ), + # --- SAM3 text-prompt inference --------------------------------------------------------- + BatchGenerate( + steps.list_sam_config, + outputs=[data.sam_config_tbl], + labels=[("stage", "prompt")], + ), + BatchTransform( + func=steps.sam_inference, + inputs=[data.local_images_tbl, data.sam_config_tbl], + outputs=[data.sam_predictions_tbl], + transform_keys=["image_id", "config_id"], + chunk_size=1, + labels=[("stage", "sam")], + executor_config=ExecutorConfig(parallelism=0), + ), + BatchTransform( + func=steps.sam_to_cvat_xml, + inputs=[data.sam_predictions_tbl], + outputs=[data.sam_cvat_xml_tbl], + transform_keys=["image_id"], + labels=[("stage", "sam")], + ), + # --- CVAT upload + sync-back ------------------------------------------------------------ + BatchTransform( + func=steps.prepare_cvat_input, + inputs=[data.local_images_tbl, data.sam_cvat_xml_tbl], + outputs=[data.image_tbl], + transform_keys=["image_id"], + labels=[("stage", "cvat")], + ), + CVATStep( + input=data.image_tbl, + output__input_batches="image_batches", + output__cvat_task="cvat_task", + output__cvat_files="cvat_images", + task_sync_table="cvat_task_sync_table", + output__cvat_annotation="cvat_annotation", + file_path_column="image_path", + labels=[("stage", "cvat")], + minimum_files_in_job=MIN_FILES_IN_JOB, + files_batch=FILES_BATCH, + segment_size=SEGMENT_SIZE, + cvat_url=CVAT_URL, + cvat_credentials=(CVAT_USERNAME, CVAT_PASSWORD), + cvat_project_id=CVAT_PROJECT_ID, + cvat_organization=CVAT_ORGANIZATION, + primary_keys=PRIMARY_KEYS, + cloud_storage_bucket=None, + delete_unannotated_tasks_only_on_update=False, + task_queue_id__name="task_queue_id", + task_name_format="[{date:%Y-%m-%d}] TaskQueue:{task_queue_id} batch:{inner_task_id}", + create_table=True, + ), + BatchTransform( + func=steps.parse_cvat_annotations, + inputs=["cvat_annotation"], + outputs=[data.image_annotations_tbl], + transform_keys=["image_id", "task_queue_id", "inner_task_id"], + labels=[("stage", "cvat")], + ), + ] +) + +ds = DataStore(DBCONN, create_meta_table=True) +# DatapipeAPI (not plain DatapipeApp) so the datapipe-app UI front (`datapipe --pipeline app api`) +# renders the pipeline graph, table browser, and per-stage run triggers. run_logs_backend defaults +# to None -> no ClickHouse needed. +app = DatapipeAPI(ds, Catalog({}), pipeline) diff --git a/examples/video_segmentation_cvat/config.py b/examples/video_segmentation_cvat/config.py new file mode 100644 index 00000000..197317f2 --- /dev/null +++ b/examples/video_segmentation_cvat/config.py @@ -0,0 +1,90 @@ +from __future__ import annotations + +import os +from pathlib import Path + +import torch +from datapipe.store.database import DBConn + +# --- video ingest / frame sampling --------------------------------------------------------------- + +# Folder you drop long egocentric videos into. Every video in here is processed. +_INPUT_VIDEO_DIR_RAW = os.environ.get("INPUT_VIDEO_DIR") +INPUT_VIDEO_DIR = Path(_INPUT_VIDEO_DIR_RAW).resolve() if _INPUT_VIDEO_DIR_RAW else None +VIDEO_SUFFIXES = {".mp4", ".mkv", ".mov", ".webm", ".m4v"} + +# Where extracted frames are written (one subfolder per video_id). +_FRAMES_DIR_RAW = os.environ.get("FRAMES_DIR") +FRAMES_DIR = ( + Path(_FRAMES_DIR_RAW).resolve() + if _FRAMES_DIR_RAW + else Path(__file__).resolve().parent / ".frames" +) + +# ffmpeg extraction rate. 1 fps gives good event coverage; 1/3fps is not meaningfully different for +# walking POV footage (see README). Dedup below removes the near-duplicates either way. +SAMPLE_FPS = float(os.environ.get("SAMPLE_FPS", "1")) + + +def _resolve_ffmpeg() -> str: + # Prefer a system ffmpeg; fall back to the binary bundled by the imageio-ffmpeg package so the + # example runs on hosts without ffmpeg on PATH. + import shutil + + exe = os.environ.get("FFMPEG_BIN") or shutil.which("ffmpeg") + if exe: + return exe + try: + import imageio_ffmpeg + + return imageio_ffmpeg.get_ffmpeg_exe() + except Exception: + return "ffmpeg" + + +FFMPEG_BIN = _resolve_ffmpeg() + +# Perceptual-hash near-duplicate threshold (Hamming distance between consecutive frames). A frame is +# kept only if it differs from the last kept frame by more than this. Higher = more aggressive dedup. +PHASH_MAX_DISTANCE = int(os.environ.get("PHASH_MAX_DISTANCE", "10")) +PHASH_SIZE = int(os.environ.get("PHASH_SIZE", "8")) # phash hash_size (bits per side) + +IMAGE_SUFFIXES = {".jpg", ".jpeg", ".png"} + +# --- SAM3 ----------------------------------------------------------------------------------------- + +DEVICE = "cuda:0" if torch.cuda.is_available() else "cpu" + +HF_TOKEN = os.environ.get("HF_TOKEN", "") +SAM_TEXT_PROMPT = os.environ.get("SAM_TEXT_PROMPT", "person") +SAM_SCORE_THRESHOLD = float(os.environ.get("SAM_SCORE_THRESHOLD", "0.5")) +SAM_MAX_DETECTIONS = int(os.environ.get("SAM_MAX_DETECTIONS", "50")) + +# --- CVAT ----------------------------------------------------------------------------------------- + +TASK_QUEUE_ID = os.environ.get("TASK_QUEUE_ID", "queue1") +# One CVAT task per video: task_queue_id is set to video_id (see steps.prepare_cvat_input), and +# FILES_BATCH is large so all of a video's frames land in a single task. +FILES_BATCH = int(os.environ.get("FILES_BATCH", "100000")) +# Minimum frames to open a task — 1 so short videos still get a task. +MIN_FILES_IN_JOB = int(os.environ.get("MIN_FILES_IN_JOB", "1")) +# CVAT jobs per task: each task is split into jobs of this many frames (None = single job). +_SEGMENT_SIZE_RAW = os.environ.get("SEGMENT_SIZE", "200") +SEGMENT_SIZE = int(_SEGMENT_SIZE_RAW) if _SEGMENT_SIZE_RAW.strip() else None + +CVAT_URL = os.environ.get("CVAT_URL", "http://localhost:8080") +CVAT_USERNAME = os.environ.get("CVAT_USERNAME", "admin") +CVAT_PASSWORD = os.environ.get("CVAT_PASSWORD", "admin") +CVAT_PROJECT_ID = int(os.environ.get("CVAT_PROJECT_ID", "1")) +CVAT_ORGANIZATION = os.environ.get("CVAT_ORGANIZATION", "") +CVAT_PROJECT_NAME = os.environ.get("CVAT_PROJECT_NAME", "datapipe-robots-ego-cvat") + +CVAT_BOX_LABEL = os.environ.get("CVAT_BOX_LABEL", "person_box") +CVAT_POLYGON_LABEL = os.environ.get("CVAT_POLYGON_LABEL", "person_mask") + +PRIMARY_KEYS = ["image_id", "task_queue_id"] + +# --- datapipe ------------------------------------------------------------------------------------- + +DB_URL = os.environ.get("DB_URL") +DBCONN = DBConn(DB_URL, None) diff --git a/examples/video_segmentation_cvat/data.py b/examples/video_segmentation_cvat/data.py new file mode 100644 index 00000000..62e010ed --- /dev/null +++ b/examples/video_segmentation_cvat/data.py @@ -0,0 +1,133 @@ +from __future__ import annotations + +from datapipe.compute import Table +from datapipe.store.database import TableStoreDB +from sqlalchemy import Column, Float, Integer, JSON, String + +from config import DBCONN + +# One row per video file found in INPUT_VIDEO_DIR. +video_tbl = Table( + name="video", + store=TableStoreDB( + dbconn=DBCONN, + name="video", + data_sql_schema=[ + Column("video_id", String, primary_key=True), + Column("video_path", String), + ], + create_table=True, + ), +) + +# One row per frame extracted from a video (before dedup). video_id is part of the PK so the +# per-video extract/dedup transforms (transform_keys=["video_id"]) can explode/reduce cleanly. +frames_tbl = Table( + name="frames", + store=TableStoreDB( + dbconn=DBCONN, + name="frames", + data_sql_schema=[ + Column("video_id", String, primary_key=True), + Column("frame_id", String, primary_key=True), + Column("ts_sec", Float), + Column("frame_path", String), + ], + create_table=True, + ), +) + +# Survivors of perceptual-hash dedup — this is what the SAM->CVAT tail consumes, so it stays identical +# to sam_cvat's `local_images`. video_id lingers in the PK for the per-video CVAT task split (and so +# dedup, grouped per video, can delete+reinsert a video's survivors); it is reduced away by sam_inference. +local_images_tbl = Table( + name="local_images", + store=TableStoreDB( + dbconn=DBCONN, + name="local_images", + data_sql_schema=[ + Column("video_id", String, primary_key=True), + Column("image_id", String, primary_key=True), + Column("image_path", String), + ], + create_table=True, + ), +) + +sam_config_tbl = Table( + name="sam_config", + store=TableStoreDB( + dbconn=DBCONN, + name="sam_config", + data_sql_schema=[ + Column("config_id", String, primary_key=True), + Column("text_prompt", String), + ], + create_table=True, + ), +) + +sam_predictions_tbl = Table( + name="sam_predictions", + store=TableStoreDB( + dbconn=DBCONN, + name="sam_predictions", + data_sql_schema=[ + Column("image_id", String, primary_key=True), + Column("detection_id", String, primary_key=True), + Column("score", Float), + Column("x_min", Float), + Column("y_min", Float), + Column("x_max", Float), + Column("y_max", Float), + Column("polygon_points", JSON), + ], + create_table=True, + ), +) + +sam_cvat_xml_tbl = Table( + name="sam_cvat_xml", + store=TableStoreDB( + dbconn=DBCONN, + name="sam_cvat_xml", + data_sql_schema=[ + Column("image_id", String, primary_key=True), + Column("annotations", String), + ], + create_table=True, + ), +) + +image_tbl = Table( + name="image", + store=TableStoreDB( + dbconn=DBCONN, + name="image", + data_sql_schema=[ + Column("image_id", String, primary_key=True), + Column("task_queue_id", String, primary_key=True), + Column("image_path", String), + Column("annotations", String), + ], + create_table=True, + ), +) + +image_annotations_tbl = Table( + name="image__annotations", + store=TableStoreDB( + dbconn=DBCONN, + name="image__annotations", + data_sql_schema=[ + Column("image_id", String, primary_key=True), + Column("task_queue_id", String, primary_key=True), + Column("inner_task_id", Integer, primary_key=True), + Column("boxes", JSON), + Column("polygons", JSON), + Column("box_labels", JSON), + Column("polygon_labels", JSON), + ], + create_table=True, + ), +) diff --git a/examples/video_segmentation_cvat/models.py b/examples/video_segmentation_cvat/models.py new file mode 100644 index 00000000..f5465e4a --- /dev/null +++ b/examples/video_segmentation_cvat/models.py @@ -0,0 +1,120 @@ +from __future__ import annotations + +import logging +from typing import List, Optional + +import cv2 +import numpy as np +import torch +from cv_pipeliner import BboxData +from PIL import Image + +from config import DEVICE, HF_TOKEN, SAM_MAX_DETECTIONS, SAM_SCORE_THRESHOLD + +logger = logging.getLogger(__name__) + +_processor = None + + +def ensure_hf_login() -> None: + from huggingface_hub import login + + if HF_TOKEN: + login(token=HF_TOKEN) + return + login() + + +def get_processor(): + global _processor + if _processor is None: + ensure_hf_login() + from sam3.model.sam3_image_processor import Sam3Processor + from sam3.model_builder import build_sam3_image_model + + model = build_sam3_image_model() + _processor = Sam3Processor(model) + logger.info("Loaded SAM3 model on device %s", DEVICE) + return _processor + + +def _mask_to_polygon(mask: np.ndarray) -> Optional[np.ndarray]: + mask_uint8 = (mask.astype(np.uint8) * 255) if mask.max() <= 1 else mask.astype(np.uint8) + contours, _ = cv2.findContours(mask_uint8, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) + if not contours: + return None + + contour = max(contours, key=cv2.contourArea) + epsilon = 0.005 * cv2.arcLength(contour, True) + approx = cv2.approxPolyDP(contour, epsilon, True) + if approx.shape[0] < 3: + return None + return approx.reshape(-1, 2).astype(np.int32) + + +def _to_numpy(value) -> np.ndarray: + return np.asarray(value.detach().cpu().float().numpy()) + + +def _to_scalar(value) -> float: + array = _to_numpy(value) + return float(array.reshape(-1)[0]) + + +def infer_image(image: Image.Image, text_prompt: str) -> List[BboxData]: + processor = get_processor() + device_type = "cuda" if DEVICE.startswith("cuda") else "cpu" + + if device_type == "cuda": + with torch.autocast(device_type, dtype=torch.bfloat16): + inference_state = processor.set_image(image) + output = processor.set_text_prompt(state=inference_state, prompt=text_prompt) + else: + inference_state = processor.set_image(image) + output = processor.set_text_prompt(state=inference_state, prompt=text_prompt) + + masks = output.get("masks", []) + boxes = output.get("boxes", []) + scores = output.get("scores", []) + + if masks is None or boxes is None or scores is None: + return [] + + num_detections = min(len(masks), len(boxes), len(scores)) + detections: List[BboxData] = [] + + indexed_scores = [(idx, _to_scalar(scores[idx])) for idx in range(num_detections)] + indexed_scores.sort(key=lambda item: item[1], reverse=True) + + kept = 0 + for rank, (idx, score) in enumerate(indexed_scores): + if score < SAM_SCORE_THRESHOLD: + continue + if kept >= SAM_MAX_DETECTIONS: + break + + box = _to_numpy(boxes[idx]).reshape(-1) + if box.size < 4: + continue + + x_min, y_min, x_max, y_max = [float(v) for v in box[:4]] + mask = _to_numpy(masks[idx]) + if mask.ndim == 3: + mask = mask[0] + polygon = _mask_to_polygon(mask) + mask_polygons = [polygon] if polygon is not None else [] + + detections.append( + BboxData( + xmin=x_min, + ymin=y_min, + xmax=x_max, + ymax=y_max, + detection_score=score, + mask=mask_polygons, + additional_info={"detection_id": str(rank)}, + ) + ) + kept += 1 + + return detections diff --git a/examples/video_segmentation_cvat/pyproject.toml b/examples/video_segmentation_cvat/pyproject.toml new file mode 100644 index 00000000..d9c618ac --- /dev/null +++ b/examples/video_segmentation_cvat/pyproject.toml @@ -0,0 +1,38 @@ +[project] +name = "robots-ego-cvat-example" +version = "0" +requires-python = ">=3.10,<3.13" +dependencies = [ + "torch==2.6.0", + "torchvision==0.21.0", + "setuptools<82", + "sam3", + "datapipe-core", + "datapipe-cvat", + "datapipe-app[ui]", + "datapipe-ui", + "python-dotenv", + "opencv-python", + "huggingface-hub==1.20.1", + "cv-pipeliner", + "einops", + "pycocotools", + "imagehash", + "imageio-ffmpeg", + "psutil" +] + +[[tool.uv.index]] +name = "pytorch-cu124" +url = "https://download.pytorch.org/whl/cu124" +explicit = true + +[tool.uv.sources] +torch = { index = "pytorch-cu124" } +torchvision = { index = "pytorch-cu124" } +datapipe-core = { path = "../../libs/datapipe-core", editable = true } +datapipe-cvat = { path = "../../libs/datapipe-cvat", editable = true } +datapipe-app = { path = "../../libs/datapipe-app", editable = true } +datapipe-ui = { path = "../../libs/datapipe-ui", editable = true } +cv-pipeliner = { git = "https://github.com/epoch8/cv-pipeliner", rev = "5724f8d54e4df64013fad85d41129799bc143293" } +sam3 = { git = "https://github.com/facebookresearch/sam3.git", rev = "5dd401d1c5c1d5c3eedff06d41b77af824517619" } diff --git a/examples/video_segmentation_cvat/scripts/fetch_video.py b/examples/video_segmentation_cvat/scripts/fetch_video.py new file mode 100644 index 00000000..816f8729 --- /dev/null +++ b/examples/video_segmentation_cvat/scripts/fetch_video.py @@ -0,0 +1,78 @@ +#!/usr/bin/env python3 +"""Download the built-in egocentric test videos into INPUT_VIDEO_DIR from the demo bucket. + +The ~24h city-walk set (Japan ~8h, USA ~8h, India ~8h across 5 clips) plus a 3-min smoke clip live in +an object-storage bucket as pre-encoded 720p `.webm`. Override the base URL with $VIDEO_BUCKET_URL. + +Internal-demo use only: the source footage is under YouTube ToS — do not redistribute the videos or +the frames sampled from them. + +Usage: + python scripts/fetch_video.py # the whole built-in test set + python scripts/fetch_video.py smoke_shibuya_3min # just the 3-min smoke clip (live demo) + python scripts/fetch_video.py 27Pv4Cg4EV4 BsiHD4m6_BU # specific bucket keys +""" +from __future__ import annotations + +import argparse +import os +import subprocess +import sys +from pathlib import Path + +# Base URL of the demo bucket holding the built-in test videos (public-read). Override via env. +BUCKET_URL = os.environ.get( + "VIDEO_BUCKET_URL", "https://storage.yandexcloud.net/e8-demo/robots-ego-video" +) + +# Built-in test set as bucket keys (without .webm). The 3-min smoke clip is listed separately so it +# can be fetched on its own for the live-demo stage (see the setup skill's two-stage choreography). +CITY_WALKS = [ + "BsiHD4m6_BU", # Tokyo, 9 districts, ~8h + "27Pv4Cg4EV4", # New York full city walk, ~8h + "60Q5E0KZb38", # Mumbai markets, ~2h40 + "qskdzPj39hE", # New Delhi Paharganj, ~2h + "8W4ZTX1z02E", # Mumbai busy streets, ~1h35 + "7wBNtsgqNOI", # New Delhi crowds, ~1h + "Lteooc0BHtk", # New Delhi streets, ~40m +] +SMOKE = "smoke_shibuya_3min" # 3-min live-demo clip +DEFAULT_KEYS = CITY_WALKS + [SMOKE] + + +def main() -> int: + parser = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) + parser.add_argument( + "keys", nargs="*", help="bucket keys to fetch, without .webm (default: the whole built-in set)" + ) + parser.add_argument( + "--dir", default=os.environ.get("INPUT_VIDEO_DIR"), help="target dir (default: $INPUT_VIDEO_DIR)" + ) + parser.add_argument( + "--base-url", default=BUCKET_URL, help="bucket base URL (default: $VIDEO_BUCKET_URL)" + ) + args = parser.parse_args() + + if not args.dir: + parser.error("set --dir or the INPUT_VIDEO_DIR env var") + out_dir = Path(args.dir).resolve() + out_dir.mkdir(parents=True, exist_ok=True) + + keys = args.keys or DEFAULT_KEYS + base = args.base_url.rstrip("/") + rc = 0 + for key in keys: + key = key[:-5] if key.endswith(".webm") else key + url = f"{base}/{key}.webm" + dst = out_dir / f"{key}.webm" + print(f"-> {dst} <- {url}", file=sys.stderr) + # -f fail on HTTP errors, -S show errors, -L follow redirects, -C - resume a partial download. + ret = subprocess.run(["curl", "-fSL", "-C", "-", "-o", str(dst), url]).returncode + rc = rc or ret + return rc + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/examples/video_segmentation_cvat/steps.py b/examples/video_segmentation_cvat/steps.py new file mode 100644 index 00000000..bfc03c99 --- /dev/null +++ b/examples/video_segmentation_cvat/steps.py @@ -0,0 +1,315 @@ +from __future__ import annotations + +import subprocess +import xml.etree.ElementTree as ET +from pathlib import Path +from typing import Iterator + +import imagehash +import pandas as pd +from cv_pipeliner import BboxData +from PIL import Image + +from config import ( + CVAT_BOX_LABEL, + CVAT_POLYGON_LABEL, + FFMPEG_BIN, + FRAMES_DIR, + INPUT_VIDEO_DIR, + PHASH_MAX_DISTANCE, + PHASH_SIZE, + SAM_TEXT_PROMPT, + SAMPLE_FPS, + TASK_QUEUE_ID, + VIDEO_SUFFIXES, +) +from models import infer_image + +SAM_CONFIG_ID = "default" + +_LOCAL_IMAGES_COLUMNS = ["video_id", "image_id", "image_path"] +_FRAMES_COLUMNS = ["video_id", "frame_id", "ts_sec", "frame_path"] +_SAM_PRED_COLUMNS = [ + "image_id", + "detection_id", + "score", + "x_min", + "y_min", + "x_max", + "y_max", + "polygon_points", +] + + +# --- video -> frames ----------------------------------------------------------------------------- + + +def list_videos() -> Iterator[pd.DataFrame]: + records: list[dict[str, str]] = [] + if INPUT_VIDEO_DIR is not None and INPUT_VIDEO_DIR.exists(): + for path in sorted(INPUT_VIDEO_DIR.iterdir()): + if path.is_file() and path.suffix.lower() in VIDEO_SUFFIXES: + records.append({"video_id": path.stem, "video_path": str(path)}) + + if records: + yield pd.DataFrame(records) + else: + yield pd.DataFrame(columns=["video_id", "video_path"]) + + +def _extract_one_video(video_id: str, video_path: str) -> list[dict]: + out_dir = FRAMES_DIR / video_id + out_dir.mkdir(parents=True, exist_ok=True) + + existing = sorted(out_dir.glob("f_*.jpg")) + if not existing: + # -vf fps=N samples N frames per second; VFR keeps timestamps honest. + cmd = [ + FFMPEG_BIN, + "-hide_banner", + "-loglevel", + "error", + "-i", + video_path, + "-vf", + f"fps={SAMPLE_FPS}", + "-q:v", + "2", + str(out_dir / "f_%06d.jpg"), + ] + subprocess.run(cmd, check=True) + existing = sorted(out_dir.glob("f_*.jpg")) + + records = [] + for path in existing: + n = int(path.stem.split("_")[-1]) # f_000123 -> 123 (1-based) + records.append( + { + "frame_id": f"{video_id}_{n:06d}", + "video_id": video_id, + "ts_sec": (n - 1) / SAMPLE_FPS, + "frame_path": str(path), + } + ) + return records + + +def extract_frames(df_video: pd.DataFrame) -> pd.DataFrame: + if df_video.empty: + return pd.DataFrame(columns=_FRAMES_COLUMNS) + + records: list[dict] = [] + for _, row in df_video.iterrows(): + records.extend(_extract_one_video(str(row["video_id"]), str(row["video_path"]))) + + if not records: + return pd.DataFrame(columns=_FRAMES_COLUMNS) + return pd.DataFrame(records, columns=_FRAMES_COLUMNS) + + +def dedup_frames(df_frames: pd.DataFrame) -> pd.DataFrame: + """Perceptual-hash dedup per video: keep a frame only if it differs from the last kept frame by + more than PHASH_MAX_DISTANCE (Hamming). Walking POV frames are highly redundant, so a sequential + compare-to-last-kept pass removes near-duplicates cheaply while preserving genuine scene change.""" + if df_frames.empty: + return pd.DataFrame(columns=_LOCAL_IMAGES_COLUMNS) + + records: list[dict] = [] + for video_id, group in df_frames.groupby("video_id"): + ordered = group.sort_values(["ts_sec", "frame_id"]) + last_hash = None + for _, row in ordered.iterrows(): + try: + with Image.open(row["frame_path"]) as img: + current = imagehash.phash(img.convert("RGB"), hash_size=PHASH_SIZE) + except (OSError, ValueError): + continue + if last_hash is not None and (current - last_hash) <= PHASH_MAX_DISTANCE: + continue + last_hash = current + records.append( + { + "video_id": str(video_id), + "image_id": row["frame_id"], + "image_path": row["frame_path"], + } + ) + + if not records: + return pd.DataFrame(columns=_LOCAL_IMAGES_COLUMNS) + return pd.DataFrame(records, columns=_LOCAL_IMAGES_COLUMNS) + + +# --- SAM config + inference (reused from sam_cvat) ----------------------------------------------- + + +def list_sam_config() -> Iterator[pd.DataFrame]: + yield pd.DataFrame([{"config_id": SAM_CONFIG_ID, "text_prompt": SAM_TEXT_PROMPT}]) + + +def _polygon_points_from_bbox(bbox: BboxData) -> list[list[float]] | None: + if not isinstance(bbox.mask, list) or not bbox.mask: + return None + polygon = max(bbox.mask, key=lambda points: len(points)).reshape(-1, 2) + if len(polygon) < 3: + return None + return [[float(x), float(y)] for x, y in polygon] + + +def sam_inference(df_local_images: pd.DataFrame, df_sam_config: pd.DataFrame) -> pd.DataFrame: + if df_local_images.empty: + return pd.DataFrame(columns=_SAM_PRED_COLUMNS) + + if df_sam_config.empty: + text_prompt = SAM_TEXT_PROMPT + else: + text_prompt = str(df_sam_config.iloc[0]["text_prompt"]) + + records = [] + for _, row in df_local_images.iterrows(): + image = Image.open(row["image_path"]).convert("RGB") + detections = infer_image(image, text_prompt) + for detection in detections: + records.append( + { + "image_id": row["image_id"], + "detection_id": detection.additional_info.get("detection_id", "0"), + "score": detection.detection_score, + "x_min": detection.xmin, + "y_min": detection.ymin, + "x_max": detection.xmax, + "y_max": detection.ymax, + "polygon_points": _polygon_points_from_bbox(detection), + } + ) + + if not records: + return pd.DataFrame(columns=_SAM_PRED_COLUMNS) + return pd.DataFrame(records) + + +# --- SAM predictions -> CVAT XML -> CVAT input (reused from sam_cvat) ----------------------------- + + +def _polygon_points_to_cvat_attr(points: list[list[float]]) -> str: + return ";".join(f"{x:.2f},{y:.2f}" for x, y in points) + + +def sam_to_cvat_xml(df_sam_predictions: pd.DataFrame) -> pd.DataFrame: + if df_sam_predictions.empty: + return pd.DataFrame(columns=["image_id", "annotations"]) + + records = [] + for image_id, group in df_sam_predictions.groupby("image_id"): + lines = [""] + for _, row in group.iterrows(): + lines.append( + f' ' + ) + lines.append(" ") + + polygon_points = row.get("polygon_points") + if isinstance(polygon_points, list) and len(polygon_points) >= 3: + points_attr = _polygon_points_to_cvat_attr(polygon_points) + lines.append( + f' ' + ) + lines.append(" ") + lines.append("") + records.append({"image_id": image_id, "annotations": "\n".join(lines)}) + + return pd.DataFrame(records, columns=["image_id", "annotations"]) + + +def prepare_cvat_input(df_local_images: pd.DataFrame, df_sam_cvat_xml: pd.DataFrame) -> pd.DataFrame: + if df_local_images.empty: + return pd.DataFrame(columns=["image_id", "task_queue_id", "image_path", "annotations"]) + + df = pd.merge(df_local_images, df_sam_cvat_xml, on="image_id", how="left") + # One CVAT task per video: scope the CVAT batch by video_id (falls back to a constant if absent). + df["task_queue_id"] = df["video_id"] if "video_id" in df.columns else TASK_QUEUE_ID + df["image_path"] = df["image_path"].apply(lambda path: str(Path(path))) + df["annotations"] = df["annotations"].fillna("") + return df[["image_id", "task_queue_id", "image_path", "annotations"]] + + +# --- CVAT annotations -> datapipe (reused from sam_cvat) ------------------------------------------ + + +def _parse_points_attr(points_attr: str) -> list[list[float]]: + points: list[list[float]] = [] + for pair in points_attr.split(";"): + if not pair.strip(): + continue + x_str, y_str = pair.split(",", 1) + points.append([float(x_str), float(y_str)]) + return points + + +def parse_cvat_annotations(df_cvat_annotation: pd.DataFrame) -> pd.DataFrame: + columns = [ + "image_id", + "task_queue_id", + "inner_task_id", + "boxes", + "polygons", + "box_labels", + "polygon_labels", + ] + if df_cvat_annotation.empty: + return pd.DataFrame(columns=columns) + + records = [] + for _, row in df_cvat_annotation.iterrows(): + annotation_xml = row.get("annotations") or "" + if not annotation_xml.strip(): + continue + + try: + image_element = ET.fromstring(annotation_xml) + except ET.ParseError: + continue + + boxes = [] + box_labels = [] + for box in image_element.findall("box"): + boxes.append( + { + "xtl": float(box.attrib.get("xtl", 0)), + "ytl": float(box.attrib.get("ytl", 0)), + "xbr": float(box.attrib.get("xbr", 0)), + "ybr": float(box.attrib.get("ybr", 0)), + } + ) + box_labels.append(box.attrib.get("label", "")) + + polygons = [] + polygon_labels = [] + for polygon in image_element.findall("polygon"): + points_attr = polygon.attrib.get("points", "") + if not points_attr: + continue + polygons.append(_parse_points_attr(points_attr)) + polygon_labels.append(polygon.attrib.get("label", "")) + + if not boxes and not polygons: + continue + + records.append( + { + "image_id": row["image_id"], + "task_queue_id": row["task_queue_id"], + "inner_task_id": int(row["inner_task_id"]), + "boxes": boxes, + "polygons": polygons, + "box_labels": box_labels, + "polygon_labels": polygon_labels, + } + ) + + if not records: + return pd.DataFrame(columns=columns) + return pd.DataFrame(records) diff --git a/libs/datapipe-cvat/datapipe_cvat/cvat_step.py b/libs/datapipe-cvat/datapipe_cvat/cvat_step.py index 5a98bc58..2074f1e2 100644 --- a/libs/datapipe-cvat/datapipe_cvat/cvat_step.py +++ b/libs/datapipe-cvat/datapipe_cvat/cvat_step.py @@ -444,6 +444,7 @@ def get_or_create_task( task_queue_id: Any, max_attempts: int, attempt_poll_s: int, + segment_size: Optional[int] = None, ) -> Tuple[Task, pd.DataFrame]: """ Creates a new task in CVAT or returns an existing one, and also associates it with images. @@ -509,6 +510,7 @@ def get_or_create_task( spec=TaskWriteRequest( name=new_task_name, project_id=project_id, + **({"segment_size": segment_size} if segment_size else {}), ), resources=resources, resource_type=resource_type, @@ -834,6 +836,7 @@ def upload_batches_to_cvat( task_name_format: str, max_attempts: int, attempt_poll_s: int, + segment_size: Optional[int] = None, failure_hook: Optional[CVATFailureHook] = None, ) -> Tuple[pd.DataFrame, pd.DataFrame]: """ @@ -1011,6 +1014,7 @@ def upload_batches_to_cvat( task_queue_id__name=task_queue_id__name, max_attempts=max_attempts, attempt_poll_s=attempt_poll_s, + segment_size=segment_size, ) task_record = { @@ -1140,6 +1144,7 @@ class CVATStep(PipelineStep): file_type: Literal["image", "video"] = "image" files_batch: Union[int, dict[Any, int]] = 100 minimum_files_in_job: Union[int, dict[Any, int]] = 50 + segment_size: Optional[int] = None # CVAT jobs per task: split each task into jobs of N frames task_queue_id__name: str = "task_queue_id" task_name_format: str = "[{date:%Y-%m-%d}] {task_queue_id} batch={inner_task_id}" sampling_order: Literal["default", "random"] = "default" @@ -1335,6 +1340,7 @@ def _mk(dt_name: str, schema: List[Column]): task_queue_id__name=self.task_queue_id__name, max_attempts=self.max_attempts, attempt_poll_s=self.attempt_poll_s, + segment_size=self.segment_size, ), executor_config=ExecutorConfig(parallelism=0), ),