Spaces:
Sleeping
Sleeping
| """Tracking and best-frame selection for price tags. | |
| The tracker keeps several top-quality crops per track. OCR is later run on the | |
| best K candidates, which is more robust than using a single frame selected only | |
| by bbox area. | |
| """ | |
| from __future__ import annotations | |
| import logging | |
| from dataclasses import dataclass, field | |
| import cv2 | |
| import numpy as np | |
| from shelf.detect.detector import Detection | |
| from shelf.io.video import frame_quality_score, laplacian_sharpness | |
| logger = logging.getLogger(__name__) | |
| class TrackCandidate: | |
| """One OCR candidate crop for a track.""" | |
| det: Detection | |
| crop: "np.ndarray" | |
| timestamp_ms: float | |
| score: float | |
| sharpness: float | |
| class TrackState: | |
| """Состояние одного трека (ценника).""" | |
| track_id: int | |
| best_det: Detection | |
| best_score: float = 0.0 | |
| best_frame: "np.ndarray | None" = field(default=None, repr=False) | |
| best_ts: float = 0.0 | |
| n_seen: int = 0 | |
| candidates: list[TrackCandidate] = field(default_factory=list, repr=False) | |
| # Frame with the sharpest QR zone — scored independently of OCR quality. | |
| # Used as a fallback decode attempt when standard candidates miss the QR. | |
| best_qr_frame: "np.ndarray | None" = field(default=None, repr=False) | |
| best_qr_score: float = 0.0 | |
| best_qr_ts: float = 0.0 | |
| def _qr_zone_sharpness(crop: np.ndarray) -> float: | |
| """Laplacian variance of the expected QR-code zone in a price-tag crop. | |
| In the raw (unrotated) crop the tag is mounted 90°-CW, so the QR code | |
| occupies the bottom-right quadrant. A high value means the QR zone is | |
| sharp — more likely to be decodable by WeChatQR. | |
| """ | |
| h, w = crop.shape[:2] | |
| roi = crop[int(h * 0.40) :, int(w * 0.40) :] # bottom-right quadrant | |
| if roi.size == 0: | |
| return 0.0 | |
| gray = cv2.cvtColor(roi, cv2.COLOR_BGR2GRAY) if roi.ndim == 3 else roi | |
| return float(cv2.Laplacian(gray, cv2.CV_64F).var()) | |
| def _clip_box( | |
| det: Detection, frame_shape: tuple[int, ...], margin: int = 0 | |
| ) -> tuple[int, int, int, int]: | |
| h, w = frame_shape[:2] | |
| x1 = max(0, int(det.x_min) - margin) | |
| y1 = max(0, int(det.y_min) - margin) | |
| x2 = min(w, int(det.x_max) + margin) | |
| y2 = min(h, int(det.y_max) + margin) | |
| return x1, y1, x2, y2 | |
| def _iou(a: Detection, b: Detection) -> float: | |
| ix = max(0, min(a.x_max, b.x_max) - max(a.x_min, b.x_min)) | |
| iy = max(0, min(a.y_max, b.y_max) - max(a.y_min, b.y_min)) | |
| inter = ix * iy | |
| union = a.area + b.area - inter | |
| return inter / union if union > 0 else 0.0 | |
| class Tracker: | |
| """ByteTrack wrapper with OCR-oriented best-frame selection. | |
| If supervision/ByteTrack is unavailable, a conservative IoU tracker is used | |
| as a fallback so the project remains runnable in a minimal environment. | |
| """ | |
| def __init__( | |
| self, | |
| lost_track_buffer: int = 30, | |
| min_hits: int = 2, | |
| track_thresh: float = 0.1, | |
| crop_margin: int = 20, | |
| max_candidates: int = 3, | |
| ): | |
| self._tracker = None | |
| self._fallback_tracks: dict[int, Detection] = {} | |
| self._next_id = 1 | |
| self.lost_track_buffer = lost_track_buffer | |
| self.min_hits = min_hits | |
| self.track_thresh = track_thresh | |
| self.crop_margin = crop_margin | |
| self.max_candidates = max_candidates | |
| self._states: dict[int, TrackState] = {} | |
| self._use_fallback = False | |
| def _load(self) -> None: | |
| try: | |
| import supervision as sv | |
| self._tracker = sv.ByteTrack( | |
| lost_track_buffer=self.lost_track_buffer, | |
| minimum_matching_threshold=0.3, | |
| ) | |
| self._use_fallback = False | |
| logger.info("ByteTrack инициализирован") | |
| except Exception as exc: | |
| self._tracker = None | |
| self._use_fallback = True | |
| logger.warning( | |
| "supervision.ByteTrack недоступен (%s); используем IoU fallback", | |
| exc, | |
| ) | |
| def update( | |
| self, | |
| detections: list[Detection], | |
| frame: np.ndarray, | |
| timestamp: float, | |
| crop_fn=None, | |
| ) -> list[tuple[int, Detection]]: | |
| """Обновить трекер. Вернуть [(track_id, det)] для активных треков. | |
| ``timestamp`` is milliseconds from the beginning of the video. | |
| ``crop_fn`` optional callable(frame, bbox) -> np.ndarray for custom crop | |
| extraction (e.g. undistort_crop_at_orig_bbox). When None, uses frame[y:y2, x:x2]. | |
| """ | |
| if self._tracker is None and not self._use_fallback: | |
| self._load() | |
| if self._use_fallback: | |
| tracked = self._update_iou_fallback(detections) | |
| else: | |
| tracked = self._update_bytetrack(detections) | |
| for tid, det in tracked: | |
| self._update_state(tid, det, frame, timestamp, crop_fn=crop_fn) | |
| return tracked | |
| def _update_bytetrack( | |
| self, detections: list[Detection] | |
| ) -> list[tuple[int, Detection]]: | |
| import supervision as sv | |
| if not detections: | |
| self._tracker.update_with_detections(sv.Detections.empty()) | |
| return [] | |
| boxes = np.array( | |
| [[d.x_min, d.y_min, d.x_max, d.y_max] for d in detections], | |
| dtype=np.float32, | |
| ) | |
| confs = np.array([d.confidence for d in detections], dtype=np.float32) | |
| class_ids = np.zeros(len(detections), dtype=int) | |
| sv_dets = sv.Detections( | |
| xyxy=boxes, confidence=confs, class_id=class_ids | |
| ) | |
| tracked = self._tracker.update_with_detections(sv_dets) | |
| results: list[tuple[int, Detection]] = [] | |
| if tracked.tracker_id is None: | |
| return results | |
| for i, tid in enumerate(tracked.tracker_id): | |
| x1, y1, x2, y2 = tracked.xyxy[i].astype(int) | |
| conf = ( | |
| float(tracked.confidence[i]) | |
| if tracked.confidence is not None | |
| else 0.5 | |
| ) | |
| det = Detection( | |
| x_min=x1, | |
| y_min=y1, | |
| x_max=x2, | |
| y_max=y2, | |
| confidence=conf, | |
| cls_name="tag", | |
| ) | |
| results.append((int(tid), det)) | |
| return results | |
| def _update_iou_fallback( | |
| self, detections: list[Detection] | |
| ) -> list[tuple[int, Detection]]: | |
| """Simple tracker for environments without supervision.""" | |
| assigned: set[int] = set() | |
| results: list[tuple[int, Detection]] = [] | |
| new_tracks: dict[int, Detection] = {} | |
| for det in detections: | |
| best_tid = None | |
| best_iou = 0.0 | |
| for tid, prev in self._fallback_tracks.items(): | |
| if tid in assigned: | |
| continue | |
| val = _iou(det, prev) | |
| if val > best_iou: | |
| best_iou = val | |
| best_tid = tid | |
| if best_tid is None or best_iou < 0.25: | |
| best_tid = self._next_id | |
| self._next_id += 1 | |
| assigned.add(best_tid) | |
| new_tracks[best_tid] = det | |
| results.append((best_tid, det)) | |
| self._fallback_tracks = new_tracks | |
| return results | |
| def _update_state( | |
| self, | |
| tid: int, | |
| det: Detection, | |
| frame: np.ndarray, | |
| timestamp_ms: float, | |
| crop_fn=None, | |
| ) -> None: | |
| x1, y1, x2, y2 = _clip_box(det, frame.shape, self.crop_margin) | |
| if crop_fn is not None: | |
| crop = crop_fn(frame, (x1, y1, x2, y2)) | |
| else: | |
| crop = frame[y1:y2, x1:x2] | |
| if crop.size == 0: | |
| return | |
| sharp = laplacian_sharpness(crop) | |
| # Larger crop area helps OCR, but glare/blur are penalized by frame_quality_score. | |
| score = float(det.area) * (1.0 + frame_quality_score(crop)) | |
| qr_score = _qr_zone_sharpness(crop) | |
| if tid not in self._states: | |
| self._states[tid] = TrackState( | |
| track_id=int(tid), | |
| best_det=det, | |
| best_score=score, | |
| best_frame=crop.copy(), | |
| best_ts=timestamp_ms, | |
| n_seen=1, | |
| best_qr_frame=crop.copy(), | |
| best_qr_score=qr_score, | |
| best_qr_ts=timestamp_ms, | |
| ) | |
| else: | |
| state = self._states[tid] | |
| state.n_seen += 1 | |
| if score > state.best_score: | |
| state.best_det = det | |
| state.best_score = score | |
| state.best_frame = crop.copy() | |
| state.best_ts = timestamp_ms | |
| if qr_score > state.best_qr_score: | |
| state.best_qr_score = qr_score | |
| state.best_qr_frame = crop.copy() | |
| state.best_qr_ts = timestamp_ms | |
| cand = TrackCandidate( | |
| det=det, | |
| crop=crop.copy(), | |
| timestamp_ms=timestamp_ms, | |
| score=score, | |
| sharpness=sharp, | |
| ) | |
| state = self._states[tid] | |
| state.candidates.append(cand) | |
| state.candidates.sort(key=lambda c: c.score, reverse=True) | |
| del state.candidates[self.max_candidates :] | |
| def get_best_crops( | |
| self, min_hits: int | None = None | |
| ) -> dict[int, TrackState]: | |
| """Вернуть лучшие кадры для треков с min_hits+ активных кадров.""" | |
| threshold = min_hits if min_hits is not None else self.min_hits | |
| return { | |
| tid: s for tid, s in self._states.items() if s.n_seen >= threshold | |
| } | |
| def reset(self) -> None: | |
| self._states.clear() | |
| self._fallback_tracks.clear() | |
| self._next_id = 1 | |
| if self._tracker is not None: | |
| try: | |
| import supervision as sv | |
| self._tracker = sv.ByteTrack( | |
| lost_track_buffer=self.lost_track_buffer, | |
| minimum_matching_threshold=0.3, | |
| ) | |
| except Exception: | |
| self._tracker = None | |
| self._use_fallback = True | |