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

"""
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)
        identity_level = cls_result.get("identity_level", "rejected")
        best_sku = cls_result.get("best_sku")
        product_info = self.inventory.get_product(best_sku) if best_sku else None

        vis_match = None
        if best_sku:
            vis_match = self.matcher.generate_match_visualization(crop, best_sku, method=method)

        annotated_shelf = self.detector.draw_detections(
            self.current_image, self.current_detections, selected_box_id=target_box["id"]
        )

        click_status = "success" if identity_level == "accepted" else identity_level
        return {
            "status": click_status,
            "clicked_point": (x, y),
            "target_box": target_box,
            "classification": {
                "method": method,
                "matched_sku": best_sku,
                "identity_level": identity_level,
                "confidence": cls_result["confidence"],
                "score": cls_result["score"],
                "top_candidates": cls_result["ranked_candidates"],
                "message": cls_result.get("message"),
            },
            "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']}")
        cls = res.get("classification") or {}
        print(f"Identity: {cls.get('identity_level')} SKU={cls.get('matched_sku')} score={cls.get('score')}")
        prod = res.get("inventory")
        if prod:
            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']}")