129 lines
4.1 KiB
Python
129 lines
4.1 KiB
Python
"""
|
|
SIFT / ORB matching: one query image vs the fixed template set.
|
|
Template keypoints/descriptors are precomputed once at startup and reused
|
|
for every request -- only the query image is processed per request.
|
|
"""
|
|
|
|
import cv2
|
|
import numpy as np
|
|
|
|
import config
|
|
|
|
_sift = None
|
|
_orb = None
|
|
# method -> {template_name: [{"kp","des"}, ...]} -- index 0 is always the
|
|
# canonical template; any further entries are hidden "hard example" photos
|
|
# (see config.HARD_EXAMPLES_DIR) matched too but never shown in the UI.
|
|
_template_features = {"SIFT": None, "ORB": None}
|
|
|
|
|
|
def _get_sift():
|
|
global _sift
|
|
if _sift is None:
|
|
_sift = cv2.SIFT_create()
|
|
return _sift
|
|
|
|
|
|
def _get_orb():
|
|
global _orb
|
|
if _orb is None:
|
|
_orb = cv2.ORB_create(nfeatures=config.ORB_N_FEATURES)
|
|
return _orb
|
|
|
|
|
|
def _detector_for(method):
|
|
return _get_sift() if method == "SIFT" else _get_orb()
|
|
|
|
|
|
def _matcher_for(method):
|
|
if method == "SIFT":
|
|
index_params = dict(algorithm=1, trees=5) # FLANN_INDEX_KDTREE
|
|
else:
|
|
index_params = dict(algorithm=6, table_number=6, key_size=12,
|
|
multi_probe_level=1) # FLANN_INDEX_LSH
|
|
search_params = dict(checks=50)
|
|
return cv2.FlannBasedMatcher(index_params, search_params)
|
|
|
|
|
|
def extract_features(method, bgr, mask):
|
|
gray = cv2.cvtColor(bgr, cv2.COLOR_BGR2GRAY)
|
|
detector = _detector_for(method)
|
|
kp, des = detector.detectAndCompute(gray, mask)
|
|
return kp, des
|
|
|
|
|
|
def set_template_features(method, template_name, kp, des):
|
|
if _template_features[method] is None:
|
|
_template_features[method] = {}
|
|
_template_features[method][template_name] = [{"kp": kp, "des": des}]
|
|
|
|
|
|
def add_hard_example_features(method, template_name, kp, des):
|
|
"""Appends an extra ("hard example") gallery entry for a template that
|
|
already has its canonical entry set. Never call before
|
|
set_template_features for the same template_name."""
|
|
_template_features[method][template_name].append({"kp": kp, "des": des})
|
|
|
|
|
|
def get_template_features(method):
|
|
return _template_features[method] or {}
|
|
|
|
|
|
def match_pair(method, kp_q, des_q, kp_t, des_t):
|
|
"""Returns (inlier_count, confidence_pct)."""
|
|
if des_q is None or des_t is None or len(kp_q) == 0 or len(kp_t) == 0:
|
|
return 0, 0.0
|
|
|
|
matcher = _matcher_for(method)
|
|
try:
|
|
raw_matches = matcher.knnMatch(des_q, des_t, k=2)
|
|
except cv2.error:
|
|
return 0, 0.0
|
|
|
|
good = []
|
|
for pair in raw_matches:
|
|
if len(pair) != 2:
|
|
continue
|
|
m, n = pair
|
|
if m.distance < config.LOWE_RATIO * n.distance:
|
|
good.append(m)
|
|
|
|
if len(good) < config.MIN_RAW_MATCHES:
|
|
return 0, 0.0
|
|
|
|
src = np.float32([kp_q[m.queryIdx].pt for m in good]).reshape(-1, 1, 2)
|
|
dst = np.float32([kp_t[m.trainIdx].pt for m in good]).reshape(-1, 1, 2)
|
|
|
|
_, ransac_mask = cv2.findHomography(src, dst, cv2.RANSAC, 5.0)
|
|
if ransac_mask is None:
|
|
return 0, 0.0
|
|
|
|
inlier_count = int(ransac_mask.sum())
|
|
confidence_pct = (inlier_count / len(good)) * 100 if good else 0.0
|
|
return inlier_count, confidence_pct
|
|
|
|
|
|
def match_against_templates(method, kp_q, des_q):
|
|
"""Returns a list of {"template": name, "score": int, "confidence": float},
|
|
sorted by score descending, for every precomputed template. Each
|
|
template is matched against its whole gallery (canonical photo + any
|
|
hidden hard examples) and the best-scoring entry wins -- a real photo
|
|
that looks more like one of the hard examples than the canonical shot
|
|
still gets full credit."""
|
|
results = []
|
|
for template_name, gallery in get_template_features(method).items():
|
|
best_inliers, best_confidence = 0, 0.0
|
|
for data in gallery:
|
|
inlier_count, confidence_pct = match_pair(
|
|
method, kp_q, des_q, data["kp"], data["des"]
|
|
)
|
|
if inlier_count > best_inliers:
|
|
best_inliers, best_confidence = inlier_count, confidence_pct
|
|
results.append({
|
|
"template": template_name,
|
|
"score": best_inliers,
|
|
"confidence": round(best_confidence, 1),
|
|
})
|
|
results.sort(key=lambda r: r["score"], reverse=True)
|
|
return results
|