0f1d46f9b3
The recognition server will classify crops from several clients at once. Two things blocked that: - OpenCV Feature2D objects were shared across calls, so parallel classify_crop() calls could corrupt each other. Each thread now gets its own extractor set; the template cache is untouched. - score_crop() is extracted from classify_crop() so callers that need every SKU score (live track aggregation) do not have to re-derive it from the top-5 list. classify_crop() output is byte-identical to before (verified against the pre-change module on 82 real crops x 4 methods). - ShelfProductDetector accepts device/imgsz so the server can pin inference. Dependencies: pin ultralytics to the version that wrote the checkpoint and that the session tracker targets (8.4.153), add lap so ultralytics never pip-installs it mid-request, and add websockets for the streaming endpoint. Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com>
344 lines
13 KiB
Python
344 lines
13 KiB
Python
"""
|
|
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
|
|
import threading
|
|
from typing import Dict, List, Tuple, Optional, Any
|
|
from identity import decide_identity
|
|
|
|
|
|
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
|
|
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,
|
|
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 score_crop(
|
|
self,
|
|
crop_bgr: np.ndarray,
|
|
method: str = "hybrid"
|
|
) -> Tuple[Dict[str, float], Dict[str, Dict[str, Any]]]:
|
|
"""
|
|
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 {}, {}
|
|
|
|
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 = 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 = 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 = 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:
|
|
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}"}
|
|
|
|
scores, details = self.score_crop(crop_bgr, method=method)
|
|
ranked = decide_identity(scores, details=details, method=method)
|
|
return ranked
|
|
|
|
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.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"]
|
|
|
|
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
|