""" End-to-End Supermarket Shelf Recognition Pipeline Integrates: 1. YOLO object detector for shelf product localization 2. SURF / ORB / BRISK / Hybrid local feature matcher for classification 3. Stock CSV inventory lookup for real-time pricing and stock quantities """ import os import sys import cv2 import numpy as np from typing import Dict, Optional, Any current_dir = os.path.dirname(os.path.abspath(__file__)) if current_dir not in sys.path: sys.path.insert(0, current_dir) from detector import ShelfProductDetector from feature_matcher import ProductFeatureMatcher from inventory import InventoryManager class ShelfRecognitionPipeline: """Master pipeline managing interactive supermarket shelf recognition.""" def __init__( self, weights_path: str = "runs/detect/train/weights/best.pt", stock_dir: str = "stock", stock_csv: str = "stock.csv" ): print("[Pipeline] Initializing Inventory Manager...") self.inventory = InventoryManager(stock_dir=stock_dir, csv_filename=stock_csv) print("[Pipeline] Initializing Feature Matcher...") self.matcher = ProductFeatureMatcher(stock_dir=stock_dir) print("[Pipeline] Initializing Shelf Detector...") self.detector = ShelfProductDetector(weights_path=weights_path) self.current_image_path: Optional[str] = None self.current_image: Optional[np.ndarray] = None self.current_detections: list = [] self.current_empty_slots: list = [] def set_active_shelf_image(self, image_input: Any) -> Dict[str, Any]: """Loads and runs detection on a new shelf image.""" if isinstance(image_input, str): self.current_image_path = image_input self.current_image = cv2.imread(image_input) else: self.current_image_path = "memory_buffer" self.current_image = image_input if self.current_image is None: return {"status": "error", "message": "Invalid or unreadable image"} det_result = self.detector.detect(self.current_image) self.current_detections = det_result.get("boxes", []) # Detect empty shelf spaces (out-of-stock gaps) self.current_empty_slots = self.detector.detect_empty_slots( self.current_image.shape, self.current_detections ) return { "status": "success", "image_shape": self.current_image.shape, "total_detected": len(self.current_detections), "boxes": self.current_detections, "empty_slots": self.current_empty_slots } def process_click( self, x: int, y: int, method: str = "hybrid", tolerance_px: int = 40 ) -> Dict[str, Any]: """ Handles user click at coordinate (x, y) on the shelf image: 1. Identifies the bounding box 2. Crops product ROI 3. Classifies via SURF/ORB/BRISK/Hybrid matching 4. Queries inventory for price, quantity, and metadata """ if self.current_image is None: return {"status": "error", "message": "No active shelf image loaded."} target_box = self.detector.find_box_at_point( x, y, self.current_detections, tolerance_px=tolerance_px ) if target_box is None: # Check if user clicked an empty shelf space for slot in getattr(self, "current_empty_slots", []): sx1, sy1, sx2, sy2 = slot["box"] if (sx1 - 10 <= x <= sx2 + 10) and (sy1 - 10 <= y <= sy2 + 10): return { "status": "empty_slot", "clicked_point": (x, y), "empty_slot": slot, "message": f"Chỗ trống trên kệ (Khu vực hàng hết / Row {slot['row_index']}). Ước tính thiếu ~{slot['missing_items_est']} sản phẩm." } return { "status": "not_found", "clicked_point": (x, y), "message": f"No product detected near coordinate ({x}, {y}). Please click inside a product packaging." } # Crop ROI crop = self.detector.crop_roi(self.current_image, target_box["box"]) # Classify with selected method cls_result = self.matcher.classify_crop(crop, method=method) if cls_result.get("status") != "success": return { "status": "classification_failed", "target_box": target_box, "message": cls_result.get("message", "Feature matching failed.") } best_sku = cls_result["best_sku"] product_info = self.inventory.get_product(best_sku) # Generate match visualization vis_match = self.matcher.generate_match_visualization(crop, best_sku, method=method) # Generate annotated shelf image highlighting the selected item annotated_shelf = self.detector.draw_detections( self.current_image, self.current_detections, selected_box_id=target_box["id"] ) return { "status": "success", "clicked_point": (x, y), "target_box": target_box, "classification": { "method": method, "matched_sku": best_sku, "confidence": cls_result["confidence"], "score": cls_result["score"], "top_candidates": cls_result["ranked_candidates"], }, "inventory": product_info, "crop_image": crop, "annotated_shelf": annotated_shelf, "match_visualization": vis_match } if __name__ == "__main__": import argparse parser = argparse.ArgumentParser(description="Test Shelf Recognition Pipeline") parser.add_argument("--image", type=str, default="runs/detect/train/train_batch0.jpg") parser.add_argument("--method", type=str, default="hybrid", choices=["surf", "orb", "brisk", "hybrid"]) args = parser.parse_args() pipeline = ShelfRecognitionPipeline() print(f"Loading image: {args.image}") init_res = pipeline.set_active_shelf_image(args.image) print(f"Total detected items: {init_res['total_detected']}") if init_res['boxes']: # Simulate click on the center of the first detected product first_box = init_res['boxes'][0] cx, cy = first_box['center'] print(f"Simulating click on Product #{first_box['id']} at ({cx}, {cy})...") res = pipeline.process_click(cx, cy, method=args.method) print("\n--- Recognition & Inventory Result ---") print(f"Status: {res['status']}") if res['status'] == 'success': prod = res['inventory'] print(f"Matched SKU: {res['classification']['matched_sku']} (Score: {res['classification']['score']})") print(f"Product Name: {prod['product_name']}") print(f"Price: {prod['price']} JPY") print(f"In Stock: {prod['number_on']} units (Received: {prod['number_in']}, Sold: {prod['number_out']})") print(f"Expiration Date: {prod['expiration_date']}")