product-recognition-demo/recognization/code/observe.py

"""Session-local perception and one-zone assisted scan.

Production path: observe(frame, catalog_version) returns an immutable
frame-local observation. Scan sessions own draft → finalized/cancelled
state. Detector/matcher stay replaceable adapters; tests inject fakes.

ponytail: in-memory sessions only. Durable store when a successful
experiment needs one.
"""

from __future__ import annotations

import hashlib
import threading
import time
import uuid
from dataclasses import dataclass, field
from typing import Any, Callable, Dict, List, Optional, Tuple

from identity import (
    IDENTITY_ACCEPTED,
    IDENTITY_AMBIGUOUS,
    IDENTITY_REJECTED,
    IDENTITY_UNKNOWN,
    decide_identity,
)

COORDINATE_CONVENTION = "pixel-xyxy-origin-top-left"

SESSION_DRAFT = "draft"
SESSION_FINALIZED = "finalized"
SESSION_CANCELLED = "cancelled"

CAPTURE_STILL = "still"
CAPTURE_SAMPLED = "sampled"

DEFAULT_MAX_IN_FLIGHT = 1
DEFAULT_MAX_QUEUED = 2
DEFAULT_SESSION_TTL_S = 30 * 60


class ObserveError(ValueError):
    """Caller or session contract violation."""


def content_fingerprint(frame_bytes: bytes) -> str:
    return hashlib.sha256(frame_bytes).hexdigest()


def _now() -> float:
    return time.time()


@dataclass
class FrameRequest:
    frame_id: str
    frame_bytes: bytes
    sequence: int
    captured_at: float
    session_id: str
    zone_id: str
    catalog_version: str
    capture_kind: str = CAPTURE_SAMPLED
    width: Optional[int] = None
    height: Optional[int] = None
    orientation: str = "upright"


@dataclass
class Observation:
    frame_id: str
    content_fingerprint: str
    sequence: int
    captured_at: float
    observed_at: float
    session_id: str
    zone_id: str
    catalog_version: str
    capture_kind: str
    width: int
    height: int
    orientation: str
    coordinate_convention: str
    boxes: List[Dict[str, Any]]
    identities: List[Dict[str, Any]]
    empty_slots: List[Dict[str, Any]]
    model_version: str
    matcher_method: str
    feature_extractor: str
    quality: Dict[str, Any]
    selected: bool = False

    def as_dict(self) -> Dict[str, Any]:
        return {
            "frame_id": self.frame_id,
            "content_fingerprint": self.content_fingerprint,
            "sequence": self.sequence,
            "captured_at": self.captured_at,
            "observed_at": self.observed_at,
            "session_id": self.session_id,
            "zone_id": self.zone_id,
            "catalog_version": self.catalog_version,
            "capture_kind": self.capture_kind,
            "width": self.width,
            "height": self.height,
            "orientation": self.orientation,
            "coordinate_convention": self.coordinate_convention,
            "boxes": list(self.boxes),
            "identities": list(self.identities),
            "empty_slots": list(self.empty_slots),
            "model_version": self.model_version,
            "matcher_method": self.matcher_method,
            "feature_extractor": self.feature_extractor,
            "quality": dict(self.quality),
            "selected": self.selected,
        }


class PerceptionAdapter:
    """Replaceable detector + matcher. Default is a testable fake."""

    def __init__(
        self,
        detect_fn: Optional[Callable[[bytes], Dict[str, Any]]] = None,
        classify_fn: Optional[Callable[[Dict[str, Any]], Dict[str, float]]] = None,
        model_version: str = "fake-detector",
        matcher_method: str = "hybrid",
        feature_extractor: str = "recorded",
    ):
        self._detect_fn = detect_fn
        self._classify_fn = classify_fn
        self.model_version = model_version
        self.matcher_method = matcher_method
        self.feature_extractor = feature_extractor

    def detect(self, frame_bytes: bytes) -> Dict[str, Any]:
        if self._detect_fn is not None:
            return self._detect_fn(frame_bytes)
        return {
            "width": 1,
            "height": 1,
            "boxes": [],
            "empty_slots": [],
            "quality": {"coverage": 0.0, "limitation": "no-detector"},
        }

    def classify_box(self, box: Dict[str, Any]) -> Dict[str, float]:
        if self._classify_fn is not None:
            return self._classify_fn(box)
        return {}


class PerceptionService:
    """Frame-local observe(). Shared adapter cache is allowed; results are not."""

    def __init__(self, adapter: Optional[PerceptionAdapter] = None):
        self.adapter = adapter or PerceptionAdapter()

    def observe(self, request: FrameRequest, catalog_version: str) -> Observation:
        if request.catalog_version != catalog_version:
            raise ObserveError("catalog_version mismatch")
        fingerprint = content_fingerprint(request.frame_bytes)
        detected = self.adapter.detect(request.frame_bytes)
        width = int(request.width or detected.get("width") or 0)
        height = int(request.height or detected.get("height") or 0)
        boxes = list(detected.get("boxes") or [])
        empty_slots = list(detected.get("empty_slots") or [])
        identities = []
        for box in boxes:
            scores = self.adapter.classify_box(box)
            ranked = decide_identity(scores, method=self.adapter.matcher_method)
            identities.append(
                {
                    "box_id": box.get("id"),
                    "identity_level": ranked["identity_level"],
                    "best_sku": ranked["best_sku"],
                    "score": ranked["score"],
                    "confidence": ranked["confidence"],
                    "ranked_candidates": ranked["ranked_candidates"],
                    "message": ranked["message"],
                }
            )
        quality = dict(detected.get("quality") or {})
        if not boxes:
            quality.setdefault("limitation", "no-visible-objects")
        return Observation(
            frame_id=request.frame_id,
            content_fingerprint=fingerprint,
            sequence=request.sequence,
            captured_at=request.captured_at,
            observed_at=_now(),
            session_id=request.session_id,
            zone_id=request.zone_id,
            catalog_version=catalog_version,
            capture_kind=request.capture_kind,
            width=width,
            height=height,
            orientation=request.orientation,
            coordinate_convention=COORDINATE_CONVENTION,
            boxes=boxes,
            identities=identities,
            empty_slots=empty_slots,
            model_version=self.adapter.model_version,
            matcher_method=self.adapter.matcher_method,
            feature_extractor=self.adapter.feature_extractor,
            quality=quality,
        )


@dataclass
class ScanSession:
    session_id: str
    zone_id: str
    catalog_version: str
    capture_kind: str
    created_at: float
    expires_at: float
    status: str = SESSION_DRAFT
    selected_frame_id: Optional[str] = None
    observations: Dict[str, Observation] = field(default_factory=dict)
    in_flight: int = 0
    queued: int = 0
    dropped_stale: int = 0
    review: Optional[Dict[str, Any]] = None
    finalized_at: Optional[float] = None
    cancelled_at: Optional[float] = None
    revision: int = 0
    lock: threading.Lock = field(default_factory=threading.Lock, repr=False, compare=False)

    def snapshot(self) -> Dict[str, Any]:
        selected = self.observations.get(self.selected_frame_id) if self.selected_frame_id else None
        return {
            "session_id": self.session_id,
            "zone_id": self.zone_id,
            "catalog_version": self.catalog_version,
            "capture_kind": self.capture_kind,
            "status": self.status,
            "created_at": self.created_at,
            "expires_at": self.expires_at,
            "selected_frame_id": self.selected_frame_id,
            "selected_observation": selected.as_dict() if selected else None,
            "observation_count": len(self.observations),
            "in_flight": self.in_flight,
            "queued": self.queued,
            "dropped_stale": self.dropped_stale,
            "review": dict(self.review) if self.review else None,
            "finalized_at": self.finalized_at,
            "cancelled_at": self.cancelled_at,
            "revision": self.revision,
        }


def _usable_score(observation: Observation) -> Tuple[float, float, int]:
    """Prefer more accepted identities, then fewer unknowns, then newer sequence."""
    accepted = sum(1 for i in observation.identities if i["identity_level"] == IDENTITY_ACCEPTED)
    unknown = sum(
        1
        for i in observation.identities
        if i["identity_level"] in (IDENTITY_UNKNOWN, IDENTITY_AMBIGUOUS, IDENTITY_REJECTED)
    )
    return (float(accepted), -float(unknown), int(observation.sequence))


class ScanService:
    """One-zone assisted scan: sampled frames, review, finalize, rescan."""

    def __init__(
        self,
        perception: Optional[PerceptionService] = None,
        *,
        max_in_flight: int = DEFAULT_MAX_IN_FLIGHT,
        max_queued: int = DEFAULT_MAX_QUEUED,
        ttl_s: float = DEFAULT_SESSION_TTL_S,
        clock: Callable[[], float] = _now,
    ):
        self.perception = perception or PerceptionService()
        self.max_in_flight = max_in_flight
        self.max_queued = max_queued
        self.ttl_s = ttl_s
        self._clock = clock
        self._sessions: Dict[str, ScanSession] = {}
        self._global_lock = threading.Lock()

    def start_session(
        self,
        zone_id: str,
        catalog_version: str,
        *,
        capture_kind: str = CAPTURE_SAMPLED,
        session_id: Optional[str] = None,
    ) -> ScanSession:
        now = self._clock()
        session = ScanSession(
            session_id=session_id or str(uuid.uuid4()),
            zone_id=zone_id,
            catalog_version=catalog_version,
            capture_kind=capture_kind,
            created_at=now,
            expires_at=now + self.ttl_s,
        )
        with self._global_lock:
            self._sessions[session.session_id] = session
        return session

    def get_session(self, session_id: str) -> ScanSession:
        with self._global_lock:
            session = self._sessions.get(session_id)
        if session is None:
            raise ObserveError(f"unknown session {session_id}")
        if self._clock() >= session.expires_at and session.status == SESSION_DRAFT:
            raise ObserveError("session expired")
        return session

    def observe_frame(self, request: FrameRequest) -> Observation:
        session = self.get_session(request.session_id)
        if request.zone_id != session.zone_id:
            raise ObserveError("zone mismatch")
        if request.catalog_version != session.catalog_version:
            raise ObserveError("catalog_version mismatch")
        if session.status != SESSION_DRAFT:
            raise ObserveError(f"session is {session.status}")

        fingerprint = content_fingerprint(request.frame_bytes)
        with session.lock:
            existing = session.observations.get(request.frame_id)
            if existing is not None:
                if existing.content_fingerprint != fingerprint:
                    raise ObserveError("frame_id reused with different content")
                return existing
            if session.in_flight >= self.max_in_flight:
                if session.queued >= self.max_queued:
                    session.dropped_stale += 1
                    raise ObserveError("backpressure")
                session.queued += 1
                queued = True
            else:
                session.in_flight += 1
                queued = False

        try:
            if queued:
                with session.lock:
                    session.queued -= 1
                    if session.in_flight >= self.max_in_flight:
                        session.dropped_stale += 1
                        raise ObserveError("backpressure")
                    session.in_flight += 1
            observation = self.perception.observe(request, session.catalog_version)
        finally:
            with session.lock:
                if session.in_flight > 0:
                    session.in_flight -= 1

        with session.lock:
            if session.status != SESSION_DRAFT:
                raise ObserveError(f"session is {session.status}")
            replay = session.observations.get(request.frame_id)
            if replay is not None:
                if replay.content_fingerprint != fingerprint:
                    raise ObserveError("frame_id reused with different content")
                return replay
            selected = session.observations.get(session.selected_frame_id) if session.selected_frame_id else None
            if selected is not None and request.sequence < selected.sequence:
                session.observations[request.frame_id] = observation
                session.dropped_stale += 1
                return observation
            session.observations[request.frame_id] = observation
            if selected is None or _usable_score(observation) >= _usable_score(selected):
                if selected is not None:
                    selected.selected = False
                observation.selected = True
                session.selected_frame_id = observation.frame_id
            return observation

    def review(
        self,
        session_id: str,
        *,
        actor: str,
        reason: str,
        note: Optional[str] = None,
        previous_value: Optional[Dict[str, Any]] = None,
    ) -> Dict[str, Any]:
        session = self.get_session(session_id)
        with session.lock:
            if session.status != SESSION_DRAFT:
                raise ObserveError(f"session is {session.status}")
            session.review = {
                "actor": actor,
                "reason": reason,
                "note": note,
                "previous_value": previous_value,
                "reviewed_at": self._clock(),
            }
            session.revision += 1
            return session.snapshot()

    def finalize(self, session_id: str) -> Dict[str, Any]:
        session = self.get_session(session_id)
        with session.lock:
            if session.status != SESSION_DRAFT:
                raise ObserveError(f"session is {session.status}")
            session.status = SESSION_FINALIZED
            session.finalized_at = self._clock()
            return session.snapshot()

    def cancel(self, session_id: str) -> Dict[str, Any]:
        session = self.get_session(session_id)
        with session.lock:
            if session.status == SESSION_FINALIZED:
                raise ObserveError("finalized sessions are immutable")
            session.status = SESSION_CANCELLED
            session.cancelled_at = self._clock()
            return session.snapshot()

    def rescan(self, session_id: str) -> ScanSession:
        parent = self.get_session(session_id)
        if parent.status == SESSION_DRAFT:
            raise ObserveError("cancel or finalize before rescan")
        child = self.start_session(
            parent.zone_id,
            parent.catalog_version,
            capture_kind=parent.capture_kind,
        )
        child.revision = parent.revision + 1
        child.review = {
            "actor": "system",
            "reason": "rescan",
            "previous_session_id": parent.session_id,
            "reviewed_at": self._clock(),
        }
        return child

    def compare_live_vs_still(
        self, sampled_session_id: str, still_session_id: str
    ) -> Dict[str, Any]:
        sampled = self.get_session(sampled_session_id)
        still = self.get_session(still_session_id)
        sampled_obs = sampled.observations.get(sampled.selected_frame_id) if sampled.selected_frame_id else None
        still_obs = still.observations.get(still.selected_frame_id) if still.selected_frame_id else None

        def counts(obs: Optional[Observation]) -> Dict[str, int]:
            if obs is None:
                return {"accepted": 0, "unknown": 0, "ambiguous": 0, "rejected": 0, "boxes": 0}
            levels = [i["identity_level"] for i in obs.identities]
            return {
                "accepted": levels.count(IDENTITY_ACCEPTED),
                "unknown": levels.count(IDENTITY_UNKNOWN),
                "ambiguous": levels.count(IDENTITY_AMBIGUOUS),
                "rejected": levels.count(IDENTITY_REJECTED),
                "boxes": len(obs.boxes),
            }

        return {
            "sampled": {
                "session_id": sampled.session_id,
                "capture_kind": sampled.capture_kind,
                "observation_count": len(sampled.observations),
                "dropped_stale": sampled.dropped_stale,
                "selected": counts(sampled_obs),
            },
            "still": {
                "session_id": still.session_id,
                "capture_kind": still.capture_kind,
                "observation_count": len(still.observations),
                "dropped_stale": still.dropped_stale,
                "selected": counts(still_obs),
            },
            "duplicate_frames_prevented": any(
                len({o.content_fingerprint for o in s.observations.values()}) < len(s.observations)
                for s in (sampled, still)
            ),
            "note": "Comparator of session mechanics, not measured shelf accuracy or latency.",
        }