Преглед изворни кода

feat: add multi-scale template matcher

wenhongquan пре 1 недеља
родитељ
комит
f7a5db5d06
2 измењених фајлова са 78 додато и 0 уклоњено
  1. 52 0
      calibration_scan_180_360/matchers.py
  2. 26 0
      calibration_scan_180_360/test_matchers.py

+ 52 - 0
calibration_scan_180_360/matchers.py

@@ -0,0 +1,52 @@
+from typing import Optional, Tuple
+import cv2
+import numpy as np
+
+
+class TemplateMatcher:
+    def __init__(self,
+                 roi_ratio: float = 0.5,
+                 scales: Tuple[float, ...] = (0.10, 0.12, 0.15, 0.18, 0.20, 0.25, 0.30),
+                 score_threshold: float = 0.3):
+        self.roi_ratio = roi_ratio
+        self.scales = scales
+        self.score_threshold = score_threshold
+
+    def match(self, ptz_img: np.ndarray, panorama_img: np.ndarray
+              ) -> Optional[Tuple[float, float, float, float]]:
+        if ptz_img is None or panorama_img is None:
+            return None
+
+        ph, pw = ptz_img.shape[:2]
+        rw = int(pw * self.roi_ratio)
+        rh = int(ph * self.roi_ratio)
+        x0 = (pw - rw) // 2
+        y0 = (ph - rh) // 2
+        roi = ptz_img[y0:y0+rh, x0:x0+rw]
+
+        if roi.size == 0:
+            return None
+
+        roi_gray = cv2.cvtColor(roi, cv2.COLOR_BGR2GRAY) if len(roi.shape) == 3 else roi
+        pano_gray = cv2.cvtColor(panorama_img, cv2.COLOR_BGR2GRAY) if len(panorama_img.shape) == 3 else panorama_img
+        pano_h, pano_w = pano_gray.shape
+
+        best = None
+        best_score = -1.0
+        for scale in self.scales:
+            sw = max(1, int(rw * scale))
+            sh = max(1, int(rh * scale))
+            if sw > pano_w or sh > pano_h:
+                continue
+            resized = cv2.resize(roi_gray, (sw, sh), interpolation=cv2.INTER_AREA)
+            result = cv2.matchTemplate(pano_gray, resized, cv2.TM_CCOEFF_NORMED)
+            _, max_val, _, max_loc = cv2.minMaxLoc(result)
+            if max_val > best_score:
+                best_score = max_val
+                center_x = max_loc[0] + sw / 2
+                center_y = max_loc[1] + sh / 2
+                best = (center_x / pano_w, center_y / pano_h, float(max_val), scale)
+
+        if best is None or best[2] < self.score_threshold:
+            return None
+        return best

+ 26 - 0
calibration_scan_180_360/test_matchers.py

@@ -0,0 +1,26 @@
+import numpy as np
+from matchers import TemplateMatcher
+
+
+def test_template_matcher_finds_known_location():
+    # 构造全景图 400x200,中间放一块带纹理的红色方块
+    # 注意:纯色 ROI 会让 TM_CCOEFF_NORMED 产生退化(全图相关值相同),
+    # 因此给红色通道附加单调渐变纹理,使模板匹配有唯一峰值。
+    panorama = np.zeros((200, 400, 3), dtype=np.uint8)
+    y, x = np.mgrid[80:120, 180:220]
+    panorama[80:120, 180:220, 0] = ((x - 180) * 6).astype(np.uint8)
+    panorama[80:120, 180:220, 1] = ((y - 80) * 6).astype(np.uint8)
+    panorama[80:120, 180:220, 2] = 255
+
+    # 构造 PTZ 图:只看这块红色方块的中心 ROI
+    ptz = panorama[80:120, 180:220].copy()
+
+    matcher = TemplateMatcher(scales=(0.8, 1.0, 1.2), score_threshold=0.5)
+    result = matcher.match(ptz, panorama)
+
+    assert result is not None
+    x_ratio, y_ratio, score, scale = result
+    # 中心应在 (200, 100) 附近
+    assert 0.45 <= x_ratio <= 0.55
+    assert 0.45 <= y_ratio <= 0.55
+    assert score > 0.8