feat: add intake session state machine

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

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()