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

"""
Record a live session so it can become data.

Everything measured so far comes from one or two photographs, which is the largest
weakness in the project. The fix is a store walk recorded from the device that will
do the scanning, with enough context to label it afterwards.

`--record-frames DIR` writes each processed frame plus a JSON line describing what
the server saw in it. The lines are the labelling input: a facing that the tag
channel named is a candidate label, and the confirmation screen is where a human
accepts or corrects it. Nothing here is a training set on its own - it is the raw
material for one.

Layout:

    DIR/<session>/000123.jpg          the frame, exactly as the client sent it
    DIR/<session>/index.jsonl         one line per frame:
        {"frame": 123, "path": "...", "at": 1699..., "summary": {...},
         "tracks": [{"box": [...], "identity": {...}, "tag": {...}}]}

Frames are written at most every `every` processed frames (default 3) so a long
walk does not fill a disk, and the index keeps the frame's own numbering so the
sampling is visible rather than silent.
"""

from __future__ import annotations

import json
import os
import threading
import time
from typing import Any, Dict, List, Optional


class FrameRecorder:
    """Write frames and their results to disk, for labelling later."""

    def __init__(self, root: str, every: int = 3, max_frames_per_session: int = 4000) -> None:
        self.root = root
        self.every = max(1, int(every))
        self.max_frames_per_session = int(max_frames_per_session)
        self._lock = threading.Lock()
        self._counts: Dict[str, int] = {}
        self._written: Dict[str, int] = {}
        os.makedirs(self.root, exist_ok=True)

    def _session_dir(self, session_id: str) -> str:
        path = os.path.join(self.root, session_id)
        os.makedirs(path, exist_ok=True)
        return path

    def write(self, session_id: str, frame_id: int, payload: bytes, result: Dict[str, Any]) -> None:
        """Record one processed frame and what was seen in it."""
        with self._lock:
            seen = self._counts.get(session_id, 0) + 1
            self._counts[session_id] = seen
            if seen % self.every:
                return
            written = self._written.get(session_id, 0)
            if written >= self.max_frames_per_session:
                return
            self._written[session_id] = written + 1

        directory = self._session_dir(session_id)
        name = f"{written:06d}.jpg"
        try:
            with open(os.path.join(directory, name), "wb") as handle:
                handle.write(payload)
        except OSError:
            return

        tracks: List[Dict[str, Any]] = []
        for track in result.get("tracks") or []:
            tracks.append(
                {
                    "box": track.get("box"),
                    "identity": track.get("identity"),
                    "tag_identity": track.get("tag_identity"),
                    "tag": track.get("price_tag"),
                }
            )
        line = {
            "frame": frame_id,
            "path": name,
            "at": round(time.time(), 3),
            "summary": result.get("summary"),
            "tracks": tracks,
        }
        try:
            with open(os.path.join(directory, "index.jsonl"), "a", encoding="utf-8") as handle:
                handle.write(json.dumps(line, ensure_ascii=False) + "\n")
        except OSError:
            return

    def stats(self) -> Dict[str, Any]:
        with self._lock:
            return {
                "root": self.root,
                "every": self.every,
                "sessions": len(self._counts),
                "frames_written": sum(self._written.values()),
            }


def labelling_candidates(index_path: str) -> List[Dict[str, Any]]:
    """Turn a recorded index into candidate labels for a human to confirm.

    A facing that the tag channel named carries a product name and a corroborated
    price; those are the seeds. The rest are candidates grouped by position, to be
    labelled in the confirmation screen rather than by hand from scratch.
    """
    seeds: List[Dict[str, Any]] = []
    with open(index_path, "r", encoding="utf-8") as handle:
        for line in handle:
            line = line.strip()
            if not line:
                continue
            entry = json.loads(line)
            for track in entry.get("tracks") or []:
                tag_identity = track.get("tag_identity") or {}
                tag = track.get("tag") or {}
                if not tag_identity.get("sku_id"):
                    continue
                seeds.append(
                    {
                        "frame": entry["frame"],
                        "box": track.get("box"),
                        "sku_id": tag_identity["sku_id"],
                        "name_similarity": tag_identity.get("score"),
                        "tag_price": tag.get("price"),
                        "corroborated": bool(
                            {"tax_consistent", "tax_pair_consistent", "discount_consistent"}
                            & set(tag.get("flags") or [])
                        ),
                    }
                )
    return seeds