Initial commit for push folder

recognization/code/pipeline.py

@@ -0,0 +1,181 @@
+"""
+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']}")