From a74a5d7c38897477e5ca6134b3f4cd373a4089b0 Mon Sep 17 00:00:00 2001 From: pyr0ball Date: Mon, 13 Jul 2026 10:47:47 -0700 Subject: [PATCH] feat: add intake session state machine --- bookmark/session.py | 85 ++++++++++++++++++++++++++++++++++++++ tests/test_session.py | 96 +++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 181 insertions(+) create mode 100644 bookmark/session.py create mode 100644 tests/test_session.py diff --git a/bookmark/session.py b/bookmark/session.py new file mode 100644 index 0000000..4c561c7 --- /dev/null +++ b/bookmark/session.py @@ -0,0 +1,85 @@ +from dataclasses import dataclass, field +from datetime import datetime, timezone +from enum import Enum +from pathlib import Path + +import cv2 +import numpy as np + +from .capture_strategies.base import CaptureStrategy +from .catalogue import BookRecord, Catalogue +from .image_processing import crop_to_book +from .sizing import MarkerRegion, determine_size_bucket + + +class IntakeState(str, Enum): + WAITING_FRONT = "waiting_front" + WAITING_BACK = "waiting_back" + BOOK_COMPLETE = "book_complete" + + +@dataclass +class IntakeSession: + catalogue: Catalogue + strategy: CaptureStrategy + reference_frame: np.ndarray + markers: list[MarkerRegion] + captures_dir: Path + state: IntakeState = IntakeState.WAITING_FRONT + current_book_id: str | None = None + _size_bucket: str | None = field(default=None, repr=False) + + def observe_frame(self, frame: np.ndarray) -> bool: + """Feed a frame to the active capture strategy. Returns True if a capture just fired.""" + fired = self.strategy.observe(frame) + if fired: + self._handle_capture(frame) + self.strategy.reset() + return fired + + def _handle_capture(self, frame: np.ndarray) -> None: + if self.state == IntakeState.WAITING_FRONT: + self._capture_front(frame) + elif self.state == IntakeState.WAITING_BACK: + self._capture_back(frame) + + def _capture_front(self, frame: np.ndarray) -> None: + self.current_book_id = self.catalogue.new_book_id() + self._size_bucket = determine_size_bucket(frame, self.reference_frame, self.markers) + wide_path = self.captures_dir / self.current_book_id / "front_wide.jpg" + cropped_path = self.captures_dir / self.current_book_id / "front_cropped.jpg" + self._save_pair(frame, wide_path, cropped_path) + self.catalogue.append( + BookRecord( + book_id=self.current_book_id, + timestamp=datetime.now(timezone.utc).isoformat(), + size_bucket=self._size_bucket, + front_wide_path=str(wide_path), + front_cropped_path=str(cropped_path), + ) + ) + self.state = IntakeState.WAITING_BACK + + def _capture_back(self, frame: np.ndarray) -> None: + wide_path = self.captures_dir / self.current_book_id / "back_wide.jpg" + cropped_path = self.captures_dir / self.current_book_id / "back_cropped.jpg" + self._save_pair(frame, wide_path, cropped_path) + self.catalogue.update_row( + self.current_book_id, + back_wide_path=str(wide_path), + back_cropped_path=str(cropped_path), + ) + self.state = IntakeState.BOOK_COMPLETE + + def _save_pair(self, frame: np.ndarray, wide_path: Path, cropped_path: Path) -> None: + wide_path.parent.mkdir(parents=True, exist_ok=True) + cv2.imwrite(str(wide_path), frame) + cropped = crop_to_book(frame, self.reference_frame) + cv2.imwrite(str(cropped_path), cropped) + + def start_next_book(self) -> None: + if self.state != IntakeState.BOOK_COMPLETE: + raise RuntimeError("Cannot start next book before current book is complete") + self.current_book_id = None + self._size_bucket = None + self.state = IntakeState.WAITING_FRONT diff --git a/tests/test_session.py b/tests/test_session.py new file mode 100644 index 0000000..29c6b82 --- /dev/null +++ b/tests/test_session.py @@ -0,0 +1,96 @@ +from pathlib import Path + +import numpy as np +import pytest + +from bookmark.capture_strategies.manual_button import ManualButtonStrategy +from bookmark.catalogue import Catalogue +from bookmark.session import IntakeSession, IntakeState +from bookmark.sizing import MarkerRegion + + +def make_reference(size=100, color=200): + return np.full((size, size, 3), color, dtype=np.uint8) + + +def make_book_frame(size=100, ref_color=200, book_color=50, box=(20, 20, 40, 40)): + frame = np.full((size, size, 3), ref_color, dtype=np.uint8) + x, y, w, h = box + frame[y:y + h, x:x + w] = book_color + return frame + + +@pytest.fixture +def session(tmp_path): + catalogue = Catalogue(tmp_path / "catalogue.csv") + strategy = ManualButtonStrategy() + reference = make_reference() + markers = [ + MarkerRegion(name="small_boundary", bucket="small", x=0, y=0, width=10, height=10), + MarkerRegion(name="medium_boundary", bucket="medium", x=90, y=0, width=10, height=10), + MarkerRegion(name="large_boundary", bucket="large", x=0, y=90, width=10, height=10), + ] + return IntakeSession( + catalogue=catalogue, + strategy=strategy, + reference_frame=reference, + markers=markers, + captures_dir=tmp_path / "captures", + ) + + +def test_session_starts_waiting_for_front(session): + assert session.state == IntakeState.WAITING_FRONT + + +def test_front_capture_advances_to_waiting_back_and_writes_catalogue_row(session): + session.strategy.trigger() + fired = session.observe_frame(make_book_frame()) + + assert fired is True + assert session.state == IntakeState.WAITING_BACK + assert session.current_book_id is not None + + rows = session.catalogue.read_all() + assert len(rows) == 1 + assert rows[0]["book_id"] == session.current_book_id + assert rows[0]["status"] == "captured" + assert Path(rows[0]["front_wide_path"]).exists() + assert Path(rows[0]["front_cropped_path"]).exists() + + +def test_back_capture_completes_book_and_updates_row(session): + session.strategy.trigger() + session.observe_frame(make_book_frame()) + book_id = session.current_book_id + + session.strategy.trigger() + fired = session.observe_frame(make_book_frame()) + + assert fired is True + assert session.state == IntakeState.BOOK_COMPLETE + rows = session.catalogue.read_all() + assert rows[0]["book_id"] == book_id + assert Path(rows[0]["back_wide_path"]).exists() + assert Path(rows[0]["back_cropped_path"]).exists() + + +def test_start_next_book_resets_state_and_assigns_new_id(session): + session.strategy.trigger() + session.observe_frame(make_book_frame()) + session.strategy.trigger() + session.observe_frame(make_book_frame()) + first_book_id = session.current_book_id + + session.start_next_book() + assert session.state == IntakeState.WAITING_FRONT + assert session.current_book_id is None + + session.strategy.trigger() + session.observe_frame(make_book_frame()) + assert session.current_book_id != first_book_id + + +def test_start_next_book_raises_if_book_not_complete(session): + with pytest.raises(RuntimeError): + session.start_next_book()