Add bulk frame ingest for recording sessions
Registering thousands of extracted frames one-by-one via POST /frames doesn't scale, so add `spelunkai-labeling-backend ingest-session`: it writes Frame rows directly against the database (no server needs to be running) for all frame_*.png files under a directory already placed under FRAMES_ROOT. Idempotent per (session_name, frame_index), so re-running after copying more frames only inserts the new ones. Restructured the CLI into subcommands (serve / ingest-session) while keeping `spelunkai-labeling-backend` with no arguments working exactly as before (defaults to serve), verified against a live run. 22/22 backend tests pass. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
This commit is contained in:
parent
c5fd4de09b
commit
96b7d42c34
@ -4,11 +4,13 @@ API + data model for the bounding-box labeling tool. See CLAUDE.md §3.2 for the
|
|||||||
full requirements this is built against.
|
full requirements this is built against.
|
||||||
|
|
||||||
**Status:** core data model + CRUD API implemented (sets, hierarchical classes,
|
**Status:** core data model + CRUD API implemented (sets, hierarchical classes,
|
||||||
frames, labels, per-frame/per-set status). **Not yet implemented:** dataset
|
frames, labels, per-frame/per-set status), frame image serving, and bulk frame
|
||||||
versioning/promotion (`enemy-v1`, `enemy-v2`, ...) — deferred since it needs its
|
ingest from a recording session (see `labeling/frontend/` for the UI). **Not yet
|
||||||
own design pass (snapshot semantics: labels are editable at any time, but a
|
implemented:** dataset versioning/promotion (`enemy-v1`, `enemy-v2`, ...) —
|
||||||
promoted dataset version must stay reproducible). Also not yet implemented: any
|
deferred since it needs its own design pass (snapshot semantics: labels are
|
||||||
frontend, auth/login (see below), or the active-learning auto-label workflow.
|
editable at any time, but a promoted dataset version must stay reproducible).
|
||||||
|
Also not yet implemented: auth/login (see below) or the active-learning
|
||||||
|
auto-label workflow.
|
||||||
|
|
||||||
## Stack
|
## Stack
|
||||||
|
|
||||||
@ -31,6 +33,9 @@ frontend, auth/login (see below), or the active-learning auto-label workflow.
|
|||||||
- `FrameSetStatus` — per-frame, per-set label state (`unlabeled` / `auto_labeled` /
|
- `FrameSetStatus` — per-frame, per-set label state (`unlabeled` / `auto_labeled` /
|
||||||
`reviewed`).
|
`reviewed`).
|
||||||
|
|
||||||
|
Frame images are served from `FRAMES_ROOT` (env var, default `./frames`) at
|
||||||
|
`/images/<image_path>`; `Frame.image_path` is always relative to that root.
|
||||||
|
|
||||||
## Setup
|
## Setup
|
||||||
|
|
||||||
```
|
```
|
||||||
@ -48,6 +53,24 @@ spelunkai-labeling-backend
|
|||||||
|
|
||||||
Interactive API docs at `http://127.0.0.1:8000/docs` once running.
|
Interactive API docs at `http://127.0.0.1:8000/docs` once running.
|
||||||
|
|
||||||
|
## Bulk-ingesting a recording session
|
||||||
|
|
||||||
|
After a session's extracted frames (`recording/`'s `extract-frames` output) have
|
||||||
|
been copied/rsynced onto this machine under `FRAMES_ROOT`, register them all in
|
||||||
|
one shot (writes directly to the database, no server needs to be running):
|
||||||
|
|
||||||
|
```
|
||||||
|
spelunkai-labeling-backend ingest-session \
|
||||||
|
--frames-dir /path/under/FRAMES_ROOT/run01_frames \
|
||||||
|
--session-name run01
|
||||||
|
```
|
||||||
|
|
||||||
|
`--image-path-prefix` defaults to the frames directory's own name (here
|
||||||
|
`run01_frames`) — override it if the directory was copied under a different name.
|
||||||
|
Safe to re-run: frames already registered for that session (by frame index) are
|
||||||
|
skipped. `--width`/`--height` default to 1280x720 (Spelunky Classic HD's fixed
|
||||||
|
capture resolution).
|
||||||
|
|
||||||
## Testing
|
## Testing
|
||||||
|
|
||||||
```
|
```
|
||||||
|
|||||||
@ -1,13 +1,73 @@
|
|||||||
"""Convenience entry point to run the labeling backend with uvicorn."""
|
"""Command-line entry points for the labeling backend: run the server, or
|
||||||
|
bulk-ingest a recording session's extracted frames.
|
||||||
|
"""
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import argparse
|
||||||
|
import os
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
import uvicorn
|
import uvicorn
|
||||||
|
|
||||||
|
from .db import Database
|
||||||
|
from .ingest import ingest_session
|
||||||
from .main import create_app
|
from .main import create_app
|
||||||
|
|
||||||
|
|
||||||
def main() -> None:
|
def main() -> None:
|
||||||
uvicorn.run(create_app(), host="127.0.0.1", port=8000)
|
parser = argparse.ArgumentParser(prog="spelunkai-labeling-backend")
|
||||||
|
parser.set_defaults(command="serve", host="127.0.0.1", port=8000)
|
||||||
|
subparsers = parser.add_subparsers(dest="command")
|
||||||
|
|
||||||
|
serve = subparsers.add_parser("serve", help="Run the API server (default if no subcommand is given)")
|
||||||
|
serve.add_argument("--host", default="127.0.0.1")
|
||||||
|
serve.add_argument("--port", type=int, default=8000)
|
||||||
|
|
||||||
|
ingest = subparsers.add_parser(
|
||||||
|
"ingest-session",
|
||||||
|
help="Bulk-register a recording session's extracted frames (must already be under FRAMES_ROOT)",
|
||||||
|
)
|
||||||
|
ingest.add_argument(
|
||||||
|
"--frames-dir", type=Path, required=True,
|
||||||
|
help="Directory of frame_*.png files, already placed under FRAMES_ROOT",
|
||||||
|
)
|
||||||
|
ingest.add_argument("--session-name", required=True)
|
||||||
|
ingest.add_argument(
|
||||||
|
"--image-path-prefix", default=None,
|
||||||
|
help="Path stored on each frame, relative to FRAMES_ROOT (default: --frames-dir's own name)",
|
||||||
|
)
|
||||||
|
ingest.add_argument("--width", type=int, default=1280)
|
||||||
|
ingest.add_argument("--height", type=int, default=720)
|
||||||
|
|
||||||
|
args = parser.parse_args()
|
||||||
|
|
||||||
|
if args.command == "ingest-session":
|
||||||
|
_run_ingest(args)
|
||||||
|
else:
|
||||||
|
_run_serve(args)
|
||||||
|
|
||||||
|
|
||||||
|
def _run_serve(args: argparse.Namespace) -> None:
|
||||||
|
uvicorn.run(create_app(), host=args.host, port=args.port)
|
||||||
|
|
||||||
|
|
||||||
|
def _run_ingest(args: argparse.Namespace) -> None:
|
||||||
|
database = Database(os.environ.get("DATABASE_URL"))
|
||||||
|
database.init_models()
|
||||||
|
|
||||||
|
prefix = args.image_path_prefix or args.frames_dir.name
|
||||||
|
result = ingest_session(
|
||||||
|
database=database,
|
||||||
|
frames_dir=args.frames_dir,
|
||||||
|
session_name=args.session_name,
|
||||||
|
image_path_prefix=prefix,
|
||||||
|
width=args.width,
|
||||||
|
height=args.height,
|
||||||
|
)
|
||||||
|
print(
|
||||||
|
f"Ingested {result.inserted} new frame(s) for session '{args.session_name}' "
|
||||||
|
f"({result.already_existed} already existed, {result.total_seen} total seen)."
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
if __name__ == "__main__":
|
||||||
|
|||||||
77
labeling/backend/src/spelunkai_labeling_backend/ingest.py
Normal file
77
labeling/backend/src/spelunkai_labeling_backend/ingest.py
Normal file
@ -0,0 +1,77 @@
|
|||||||
|
"""Bulk-register a recording session's extracted frames as `Frame` rows.
|
||||||
|
|
||||||
|
Meant to run on the same machine as the labeling backend's database (`jai`),
|
||||||
|
after a session's frame images have been copied/rsynced under `FRAMES_ROOT` -
|
||||||
|
inserts directly against the database rather than over HTTP, since a session
|
||||||
|
can have thousands of frames and this is a local, same-machine operation.
|
||||||
|
"""
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
from dataclasses import dataclass
|
||||||
|
from pathlib import Path
|
||||||
|
from typing import Iterator, List
|
||||||
|
|
||||||
|
from . import models
|
||||||
|
from .db import Database
|
||||||
|
|
||||||
|
FRAME_GLOB = "frame_*.png"
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass
|
||||||
|
class IngestResult:
|
||||||
|
inserted: int
|
||||||
|
already_existed: int
|
||||||
|
total_seen: int
|
||||||
|
|
||||||
|
|
||||||
|
def iter_frame_files(frames_dir: Path) -> Iterator[Path]:
|
||||||
|
return sorted(frames_dir.glob(FRAME_GLOB))
|
||||||
|
|
||||||
|
|
||||||
|
def frame_index_from_filename(path: Path) -> int:
|
||||||
|
# "frame_000123.png" -> 123, matching recording/spelunkai_recording/frames.py's naming.
|
||||||
|
return int(path.stem.split("_")[-1])
|
||||||
|
|
||||||
|
|
||||||
|
def ingest_session(
|
||||||
|
database: Database,
|
||||||
|
frames_dir: Path,
|
||||||
|
session_name: str,
|
||||||
|
image_path_prefix: str,
|
||||||
|
width: int,
|
||||||
|
height: int,
|
||||||
|
) -> IngestResult:
|
||||||
|
frame_files = list(iter_frame_files(frames_dir))
|
||||||
|
if not frame_files:
|
||||||
|
raise FileNotFoundError(f"no '{FRAME_GLOB}' files found in {frames_dir}")
|
||||||
|
|
||||||
|
db_session = database.session_factory()
|
||||||
|
try:
|
||||||
|
existing_indices = {
|
||||||
|
row.frame_index
|
||||||
|
for row in db_session.query(models.Frame.frame_index).filter_by(session_name=session_name)
|
||||||
|
}
|
||||||
|
|
||||||
|
new_frames: List[models.Frame] = []
|
||||||
|
for frame_file in frame_files:
|
||||||
|
frame_index = frame_index_from_filename(frame_file)
|
||||||
|
if frame_index in existing_indices:
|
||||||
|
continue
|
||||||
|
new_frames.append(models.Frame(
|
||||||
|
session_name=session_name,
|
||||||
|
frame_index=frame_index,
|
||||||
|
image_path=f"{image_path_prefix}/{frame_file.name}",
|
||||||
|
width=width,
|
||||||
|
height=height,
|
||||||
|
))
|
||||||
|
|
||||||
|
db_session.add_all(new_frames)
|
||||||
|
db_session.commit()
|
||||||
|
finally:
|
||||||
|
db_session.close()
|
||||||
|
|
||||||
|
return IngestResult(
|
||||||
|
inserted=len(new_frames),
|
||||||
|
already_existed=len(frame_files) - len(new_frames),
|
||||||
|
total_seen=len(frame_files),
|
||||||
|
)
|
||||||
88
labeling/backend/tests/test_ingest.py
Normal file
88
labeling/backend/tests/test_ingest.py
Normal file
@ -0,0 +1,88 @@
|
|||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
from spelunkai_labeling_backend import models
|
||||||
|
from spelunkai_labeling_backend.db import Database
|
||||||
|
from spelunkai_labeling_backend.ingest import ingest_session
|
||||||
|
|
||||||
|
|
||||||
|
def _make_frame_files(frames_dir: Path, count: int) -> None:
|
||||||
|
frames_dir.mkdir(parents=True, exist_ok=True)
|
||||||
|
for i in range(count):
|
||||||
|
(frames_dir / f"frame_{i:06d}.png").write_bytes(b"fake")
|
||||||
|
|
||||||
|
|
||||||
|
def _make_database(tmp_path: Path) -> Database:
|
||||||
|
database = Database(f"sqlite:///{tmp_path / 'test.db'}")
|
||||||
|
database.init_models()
|
||||||
|
return database
|
||||||
|
|
||||||
|
|
||||||
|
def test_ingest_session_inserts_all_frames(tmp_path):
|
||||||
|
frames_dir = tmp_path / "run01_frames"
|
||||||
|
_make_frame_files(frames_dir, 3)
|
||||||
|
database = _make_database(tmp_path)
|
||||||
|
|
||||||
|
result = ingest_session(
|
||||||
|
database=database,
|
||||||
|
frames_dir=frames_dir,
|
||||||
|
session_name="run01",
|
||||||
|
image_path_prefix="run01_frames",
|
||||||
|
width=1280,
|
||||||
|
height=720,
|
||||||
|
)
|
||||||
|
|
||||||
|
assert result.inserted == 3
|
||||||
|
assert result.already_existed == 0
|
||||||
|
assert result.total_seen == 3
|
||||||
|
|
||||||
|
db_session = database.session_factory()
|
||||||
|
try:
|
||||||
|
frames = db_session.query(models.Frame).order_by(models.Frame.frame_index).all()
|
||||||
|
assert [f.frame_index for f in frames] == [0, 1, 2]
|
||||||
|
assert [f.image_path for f in frames] == [
|
||||||
|
"run01_frames/frame_000000.png",
|
||||||
|
"run01_frames/frame_000001.png",
|
||||||
|
"run01_frames/frame_000002.png",
|
||||||
|
]
|
||||||
|
assert frames[0].width == 1280
|
||||||
|
assert frames[0].height == 720
|
||||||
|
finally:
|
||||||
|
db_session.close()
|
||||||
|
|
||||||
|
|
||||||
|
def test_ingest_session_is_idempotent_and_picks_up_new_frames(tmp_path):
|
||||||
|
frames_dir = tmp_path / "run01_frames"
|
||||||
|
_make_frame_files(frames_dir, 2)
|
||||||
|
database = _make_database(tmp_path)
|
||||||
|
|
||||||
|
ingest_session(database, frames_dir, "run01", "run01_frames", 1280, 720)
|
||||||
|
|
||||||
|
_make_frame_files(frames_dir, 3) # adds frame_000002.png
|
||||||
|
result = ingest_session(database, frames_dir, "run01", "run01_frames", 1280, 720)
|
||||||
|
|
||||||
|
assert result.inserted == 1
|
||||||
|
assert result.already_existed == 2
|
||||||
|
assert result.total_seen == 3
|
||||||
|
|
||||||
|
|
||||||
|
def test_ingest_session_keeps_sessions_independent(tmp_path):
|
||||||
|
frames_dir = tmp_path / "run01_frames"
|
||||||
|
_make_frame_files(frames_dir, 2)
|
||||||
|
database = _make_database(tmp_path)
|
||||||
|
|
||||||
|
ingest_session(database, frames_dir, "run01", "run01_frames", 1280, 720)
|
||||||
|
result = ingest_session(database, frames_dir, "run02", "run01_frames", 1280, 720)
|
||||||
|
|
||||||
|
assert result.inserted == 2
|
||||||
|
assert result.already_existed == 0
|
||||||
|
|
||||||
|
|
||||||
|
def test_ingest_session_raises_when_no_frames_found(tmp_path):
|
||||||
|
frames_dir = tmp_path / "empty"
|
||||||
|
frames_dir.mkdir()
|
||||||
|
database = _make_database(tmp_path)
|
||||||
|
|
||||||
|
with pytest.raises(FileNotFoundError):
|
||||||
|
ingest_session(database, frames_dir, "run01", "empty", 1280, 720)
|
||||||
Loading…
x
Reference in New Issue
Block a user