Initial commit for push folder

recognization/code/feature_matcher.py

@@ -0,0 +1,327 @@
+"""
+Feature Matching & Classification Module
+Implements SURF, ORB, BRISK, and Hybrid matching according to
+the methodology of Electronics 2023, 12(17), 3640.
+"""
+
+import os
+import cv2
+import numpy as np
+from typing import Dict, List, Tuple, Optional, Any
+
+
+class LocalFeatureExtractor:
+    """Encapsulates individual OpenCV local feature algorithms."""
+
+    def __init__(self, method: str = "orb"):
+        self.method = method.lower()
+        self.detector = self._create_detector(self.method)
+        self.norm_type = self._get_norm_type(self.method)
+
+    def _create_detector(self, method: str):
+        if method == "surf":
+            # Check if SURF is supported in the current OpenCV build
+            if hasattr(cv2, "xfeatures2d") and hasattr(cv2.xfeatures2d, "SURF_create"):
+                try:
+                    return cv2.xfeatures2d.SURF_create(hessianThreshold=400, nOctaves=4, nOctaveLayers=3)
+                except Exception as e:
+                    print(f"[Warning] SURF initialization failed: {e}. Falling back to SIFT.")
+            else:
+                print("[Info] SURF not available in this OpenCV build. Using SIFT as high-grade floating point descriptor.")
+            return cv2.SIFT_create(nfeatures=1000)
+
+        elif method == "orb":
+            return cv2.ORB_create(
+                nfeatures=1500,
+                scaleFactor=1.2,
+                nlevels=8,
+                edgeThreshold=31,
+                firstLevel=0,
+                WTA_K=2,
+                scoreType=cv2.ORB_HARRIS_SCORE,
+                patchSize=31,
+                fastThreshold=20
+            )
+
+        elif method == "brisk":
+            return cv2.BRISK_create(thresh=30, octaves=3, patternScale=1.0)
+
+        elif method == "sift":
+            return cv2.SIFT_create(nfeatures=1000)
+
+        else:
+            raise ValueError(f"Unsupported feature extractor method: {method}")
+
+    def _get_norm_type(self, method: str):
+        # Binary descriptors use Hamming distance; floating-point descriptors use L2 distance
+        if method in ["orb", "brisk"]:
+            return cv2.NORM_HAMMING
+        return cv2.NORM_L2
+
+    def detect_and_compute(self, image_gray: np.ndarray) -> Tuple[List[cv2.KeyPoint], Optional[np.ndarray]]:
+        """Detects keypoints and computes their descriptors."""
+        if image_gray is None or image_gray.size == 0:
+            return [], None
+        return self.detector.detectAndCompute(image_gray, None)
+
+
+class TemplateFeatureDatabase:
+    """Precomputes and caches local feature descriptors for all stock templates."""
+
+    def __init__(self, stock_dir: str = "stock"):
+        self.stock_dir = stock_dir
+        self.templates: Dict[str, Dict[str, Any]] = {}
+        self.extractors = {
+            "surf": LocalFeatureExtractor("surf"),
+            "orb": LocalFeatureExtractor("orb"),
+            "brisk": LocalFeatureExtractor("brisk")
+        }
+        self.load_and_cache_templates()
+
+    def load_and_cache_templates(self):
+        """Loads all SKU images and extracts their multi-feature descriptors."""
+        if not os.path.exists(self.stock_dir):
+            return
+
+        image_files = [
+            f for f in os.listdir(self.stock_dir)
+            if f.lower().endswith((".jpeg", ".jpg", ".png"))
+        ]
+
+        print(f"[FeatureDatabase] Pre-caching {len(image_files)} stock SKU templates...")
+        for filename in image_files:
+            sku_id = os.path.splitext(filename)[0]
+            filepath = os.path.join(self.stock_dir, filename)
+            img = cv2.imread(filepath)
+            if img is None:
+                continue
+
+            gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
+            features = {}
+            for name, extractor in self.extractors.items():
+                kp, des = extractor.detect_and_compute(gray)
+                features[name] = {"keypoints": kp, "descriptors": des}
+
+            self.templates[sku_id] = {
+                "sku_id": sku_id,
+                "filepath": filepath,
+                "image": img,
+                "gray": gray,
+                "features": features
+            }
+        print(f"[FeatureDatabase] Successfully cached {len(self.templates)} templates.")
+
+
+class ProductFeatureMatcher:
+    """
+    Main product classification engine.
+    Matches a query cropped bounding box against cached SKU templates using
+    SURF, ORB, BRISK, or the Paper's Hybrid combination.
+    """
+
+    def __init__(self, stock_dir: str = "stock"):
+        self.db = TemplateFeatureDatabase(stock_dir)
+        self.extractors = self.db.extractors
+
+    def match_descriptors(
+        self,
+        des_query: Optional[np.ndarray],
+        des_template: Optional[np.ndarray],
+        norm_type: int,
+        ratio_thresh: float = 0.75
+    ) -> List[cv2.DMatch]:
+        """Performs k-NN matching with Lowe's ratio test."""
+        if des_query is None or des_template is None:
+            return []
+        if len(des_query) < 2 or len(des_template) < 2:
+            return []
+
+        bf = cv2.BFMatcher(norm_type, crossCheck=False)
+        try:
+            raw_matches = bf.knnMatch(des_query, des_template, k=2)
+        except Exception:
+            return []
+
+        good_matches = []
+        for match_pair in raw_matches:
+            if len(match_pair) == 2:
+                m, n = match_pair
+                if m.distance < ratio_thresh * n.distance:
+                    good_matches.append(m)
+        return good_matches
+
+    def compute_geometric_inliers(
+        self,
+        kp_query: List[cv2.KeyPoint],
+        kp_template: List[cv2.KeyPoint],
+        good_matches: List[cv2.DMatch]
+    ) -> Tuple[int, List[cv2.DMatch], Optional[np.ndarray]]:
+        """Applies RANSAC Homography to verify geometric consistency."""
+        if len(good_matches) < 4:
+            return len(good_matches), good_matches, None
+
+        src_pts = np.float32([kp_query[m.queryIdx].pt for m in good_matches]).reshape(-1, 1, 2)
+        dst_pts = np.float32([kp_template[m.trainIdx].pt for m in good_matches]).reshape(-1, 1, 2)
+
+        try:
+            H, mask = cv2.findHomography(src_pts, dst_pts, cv2.RANSAC, 5.0)
+            if mask is None:
+                return len(good_matches), good_matches, None
+
+            inlier_matches = [good_matches[i] for i in range(len(good_matches)) if mask[i][0] == 1]
+            return len(inlier_matches), inlier_matches, H
+        except Exception:
+            return len(good_matches), good_matches, None
+
+    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"}
+
+        crop_gray = cv2.cvtColor(crop_bgr, cv2.COLOR_BGR2GRAY)
+        method = method.lower()
+
+        scores: Dict[str, float] = {}
+        details: Dict[str, Dict[str, Any]] = {}
+
+        if method in ["surf", "orb", "brisk"]:
+            extractor = self.extractors.get(method)
+            if extractor is None:
+                extractor = LocalFeatureExtractor(method)
+            kp_query, des_query = extractor.detect_and_compute(crop_gray)
+
+            for sku_id, tmpl in self.db.templates.items():
+                kp_tmpl = tmpl["features"][method]["keypoints"]
+                des_tmpl = tmpl["features"][method]["descriptors"]
+
+                good_matches = self.match_descriptors(
+                    des_query, des_tmpl, extractor.norm_type, ratio_thresh=0.75
+                )
+                inlier_count, inliers, _ = self.compute_geometric_inliers(
+                    kp_query, kp_tmpl, good_matches
+                )
+
+                # Score based on geometric inliers & match ratio
+                denom = max(1, min(len(kp_query), len(kp_tmpl)))
+                inlier_ratio = inlier_count / denom
+                score = inlier_count * 2.0 + inlier_ratio * 100.0
+
+                scores[sku_id] = score
+                details[sku_id] = {
+                    "inliers": inlier_count,
+                    "inlier_matches": inliers,
+                    "total_good": len(good_matches),
+                    "kp_query": kp_query,
+                    "kp_tmpl": kp_tmpl,
+                    "score": score
+                }
+
+        elif method == "hybrid":
+            # Paper's hybrid strategy combining SURF, ORB, and BRISK
+            sub_methods = ["surf", "orb", "brisk"]
+            weights = {"surf": 0.45, "orb": 0.25, "brisk": 0.30}
+
+            # 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)
+                query_features[sm] = (kp_q, des_q)
+
+            for sku_id, tmpl in self.db.templates.items():
+                composite_score = 0.0
+                sku_sub_details = {}
+
+                for sm in sub_methods:
+                    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
+
+                    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)
+
+                    denom = max(1, min(len(kp_q), len(kp_t)))
+                    sub_score = inliers * 1.5 + (inliers / denom) * 80.0
+                    composite_score += weights[sm] * sub_score
+
+                    sku_sub_details[sm] = {
+                        "inliers": inliers,
+                        "matches": len(matches),
+                        "kp_q_count": len(kp_q),
+                        "kp_t_count": len(kp_t),
+                        "sub_score": round(sub_score, 2)
+                    }
+
+                scores[sku_id] = composite_score
+                details[sku_id] = {
+                    "score": composite_score,
+                    "sub_details": sku_sub_details,
+                    "kp_query": query_features["surf"][0],
+                    "kp_tmpl": tmpl["features"]["surf"]["keypoints"]
+                }
+        else:
+            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]
+        }
+
+    def generate_match_visualization(
+        self,
+        crop_bgr: np.ndarray,
+        matched_sku: str,
+        method: str = "surf",
+        max_matches: int = 30
+    ) -> Optional[np.ndarray]:
+        """Draws visual matching lines between crop and matched template."""
+        if matched_sku not in self.db.templates:
+            return None
+
+        tmpl = self.db.templates[matched_sku]
+        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]
+        kp_q, des_q = extractor.detect_and_compute(crop_gray)
+        kp_t = tmpl["features"][eval_method]["keypoints"]
+        des_t = tmpl["features"][eval_method]["descriptors"]
+
+        matches = self.match_descriptors(des_q, des_t, extractor.norm_type, ratio_thresh=0.75)
+        _, inliers, _ = self.compute_geometric_inliers(kp_q, kp_t, matches)
+
+        matches_to_draw = inliers[:max_matches] if inliers else matches[:max_matches]
+
+        vis = cv2.drawMatches(
+            crop_bgr, kp_q,
+            tmpl["image"], kp_t,
+            matches_to_draw, None,
+            matchColor=(0, 255, 0),
+            singlePointColor=(0, 0, 255),
+            flags=cv2.DrawMatchesFlags_NOT_DRAW_SINGLE_POINTS
+        )
+        return vis