diff --git a/recognization/code/detector.py b/recognization/code/detector.py index e0062ad..cbcfada 100644 --- a/recognization/code/detector.py +++ b/recognization/code/detector.py @@ -99,8 +99,17 @@ from ultralytics import YOLO class ShelfProductDetector: """Wraps the YOLO detector for supermarket shelf products.""" - def __init__(self, weights_path: str = "runs/detect/train/weights/best.pt"): + def __init__( + self, + weights_path: str = "runs/detect/train/weights/best.pt", + device: Optional[str] = None, + imgsz: int = 640, + half: bool = False + ): self.weights_path = weights_path + self.device = device + self.imgsz = imgsz + self.half = half if not os.path.exists(weights_path): raise FileNotFoundError(f"YOLO weights not found at: {weights_path}") @@ -129,7 +138,14 @@ class ShelfProductDetector: return {"status": "error", "message": "Failed to read image", "boxes": []} h, w = image_bgr.shape[:2] - results = self.model(image_bgr, conf=conf_threshold, iou=iou_threshold, verbose=False) + predict_kwargs: Dict[str, Any] = dict( + conf=conf_threshold, iou=iou_threshold, imgsz=self.imgsz, verbose=False + ) + if self.device: + predict_kwargs["device"] = self.device + if self.half: + predict_kwargs["half"] = True + results = self.model(image_bgr, **predict_kwargs) boxes_list = [] if results and len(results) > 0: diff --git a/recognization/code/feature_matcher.py b/recognization/code/feature_matcher.py index 39f85c6..7abb325 100644 --- a/recognization/code/feature_matcher.py +++ b/recognization/code/feature_matcher.py @@ -7,7 +7,9 @@ the methodology of Electronics 2023, 12(17), 3640. import os import cv2 import numpy as np +import threading from typing import Dict, List, Tuple, Optional, Any +from identity import decide_identity class LocalFeatureExtractor: @@ -122,6 +124,18 @@ class ProductFeatureMatcher: def __init__(self, stock_dir: str = "stock"): self.db = TemplateFeatureDatabase(stock_dir) self.extractors = self.db.extractors + self._thread_local = threading.local() + + def thread_extractors(self) -> Dict[str, "LocalFeatureExtractor"]: + """One extractor set per thread: OpenCV Feature2D objects are not safe to + call concurrently from several threads.""" + extractors = getattr(self._thread_local, "extractors", None) + if extractors is None: + extractors = { + name: LocalFeatureExtractor(name) for name in ("surf", "orb", "brisk") + } + self._thread_local.extractors = extractors + return extractors def match_descriptors( self, @@ -173,26 +187,29 @@ class ProductFeatureMatcher: except Exception: return len(good_matches), good_matches, None - def classify_crop( + def score_crop( self, crop_bgr: np.ndarray, method: str = "hybrid" - ) -> Dict[str, Any]: + ) -> Tuple[Dict[str, float], Dict[str, Dict[str, Any]]]: """ - Classifies a product crop against all SKU templates. + Scores a product crop against every SKU template. + Returns (scores, details); raises ValueError for an unsupported method and + returns empty results for an empty crop. Methods supported: 'surf', 'orb', 'brisk', 'hybrid'. """ if crop_bgr is None or crop_bgr.size == 0: - return {"status": "error", "message": "Empty crop image"} + return {}, {} crop_gray = cv2.cvtColor(crop_bgr, cv2.COLOR_BGR2GRAY) method = method.lower() + extractors = self.thread_extractors() scores: Dict[str, float] = {} details: Dict[str, Dict[str, Any]] = {} if method in ["surf", "orb", "brisk"]: - extractor = self.extractors.get(method) + extractor = extractors.get(method) if extractor is None: extractor = LocalFeatureExtractor(method) kp_query, des_query = extractor.detect_and_compute(crop_gray) @@ -231,7 +248,7 @@ class ProductFeatureMatcher: # Pre-compute features for query crop across all 3 extractors query_features = {} for sm in sub_methods: - kp_q, des_q = self.extractors[sm].detect_and_compute(crop_gray) + kp_q, des_q = extractors[sm].detect_and_compute(crop_gray) query_features[sm] = (kp_q, des_q) for sku_id, tmpl in self.db.templates.items(): @@ -242,7 +259,7 @@ class ProductFeatureMatcher: kp_q, des_q = query_features[sm] kp_t = tmpl["features"][sm]["keypoints"] des_t = tmpl["features"][sm]["descriptors"] - norm_t = self.extractors[sm].norm_type + norm_t = extractors[sm].norm_type matches = self.match_descriptors(des_q, des_t, norm_t, ratio_thresh=0.75) inliers, inlier_list, _ = self.compute_geometric_inliers(kp_q, kp_t, matches) @@ -267,29 +284,29 @@ class ProductFeatureMatcher: "kp_tmpl": tmpl["features"]["surf"]["keypoints"] } else: + raise ValueError(f"Unknown method: {method}") + + return scores, details + + def classify_crop( + self, + crop_bgr: np.ndarray, + method: str = "hybrid" + ) -> Dict[str, Any]: + """ + Classifies a product crop against all SKU templates. + Methods supported: 'surf', 'orb', 'brisk', 'hybrid'. + """ + if crop_bgr is None or crop_bgr.size == 0: + return {"status": "error", "message": "Empty crop image"} + + method = method.lower() + if method not in ["surf", "orb", "brisk", "hybrid"]: return {"status": "error", "message": f"Unknown method: {method}"} - # Rank templates by score - if not scores: - return {"status": "error", "message": "No templates available"} - - sorted_candidates = sorted(scores.items(), key=lambda x: x[1], reverse=True) - top_sku, top_score = sorted_candidates[0] - runner_up_sku, runner_up_score = sorted_candidates[1] if len(sorted_candidates) > 1 else (None, 0.0) - - confidence = round(min(1.0, top_score / max(1.0, (top_score + runner_up_score * 0.5))), 4) - - return { - "status": "success", - "method": method, - "best_sku": top_sku, - "score": round(top_score, 2), - "confidence": confidence, - "ranked_candidates": [ - {"sku_id": s, "score": round(sc, 2)} for s, sc in sorted_candidates[:5] - ], - "details": details[top_sku] - } + scores, details = self.score_crop(crop_bgr, method=method) + ranked = decide_identity(scores, details=details, method=method) + return ranked def generate_match_visualization( self, @@ -306,7 +323,7 @@ class ProductFeatureMatcher: crop_gray = cv2.cvtColor(crop_bgr, cv2.COLOR_BGR2GRAY) eval_method = method if method in ["surf", "orb", "brisk"] else "surf" - extractor = self.extractors[eval_method] + extractor = self.thread_extractors()[eval_method] kp_q, des_q = extractor.detect_and_compute(crop_gray) kp_t = tmpl["features"][eval_method]["keypoints"] des_t = tmpl["features"][eval_method]["descriptors"] diff --git a/recognization/requirements-dev.txt b/recognization/requirements-dev.txt new file mode 100644 index 0000000..0d10685 --- /dev/null +++ b/recognization/requirements-dev.txt @@ -0,0 +1,3 @@ +-r requirements.txt +# FastAPI TestClient +httpx>=0.27 diff --git a/recognization/requirements.txt b/recognization/requirements.txt index 32e7a3a..809bd19 100644 --- a/recognization/requirements.txt +++ b/recognization/requirements.txt @@ -1,6 +1,11 @@ torch>=2.0.0 torchvision>=0.15.0 -ultralytics>=8.0.0 +# Exact: tracking.py relies on BYTETracker internals; the trained checkpoint was written by 8.4.152 +ultralytics==8.4.153 +# ByteTrack linear assignment; pinned so ultralytics never pip-installs it at runtime +lap>=0.5.12 +# uvicorn WebSocket support for /v1/stream +websockets>=13 opencv-python>=4.8.0 numpy>=1.23.0,<2.0.0 fastapi>=0.100.0