You cannot select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
MAI/template_matching.py

61 lines
1.8 KiB
Python

import cv2
from config import *
from helpers import clamp, box_center, box_wh
# Template matching fallback
# =========================
def _crop_safe(img, x1, y1, x2, y2):
h, w = img.shape[:2]
x1 = int(clamp(x1, 0, w - 1))
y1 = int(clamp(y1, 0, h - 1))
x2 = int(clamp(x2, 0, w))
y2 = int(clamp(y2, 0, h))
if x2 <= x1 + 1 or y2 <= y1 + 1:
return None
return img[y1:y2, x1:x2]
def tm_update_template(gray_eff, box_eff):
if gray_eff is None or box_eff is None:
return None
cx, cy = box_center(box_eff)
half = TM_TEMPLATE_SIZE // 2
patch = _crop_safe(gray_eff, cx - half, cy - half, cx + half, cy + half)
if patch is None:
return None
if patch.shape[0] < 12 or patch.shape[1] < 12:
return None
return patch.copy()
def tm_search(gray_eff, template, pred_center, miss_streak):
if (not TM_ENABLE) or gray_eff is None or template is None or pred_center is None:
return None
h, w = gray_eff.shape[:2]
cx, cy = float(pred_center[0]), float(pred_center[1])
scale = TM_SEARCH_SCALE_BASE + TM_SEARCH_SCALE_PER_MISS * float(miss_streak)
th, tw = template.shape[:2]
win_w = int(max(tw * scale, tw + 20))
win_h = int(max(th * scale, th + 20))
x1 = int(cx - win_w / 2)
y1 = int(cy - win_h / 2)
x2 = int(cx + win_w / 2)
y2 = int(cy + win_h / 2)
roi = _crop_safe(gray_eff, x1, y1, x2, y2)
if roi is None or roi.shape[0] < th + 2 or roi.shape[1] < tw + 2:
return None
res = cv2.matchTemplate(roi, template, cv2.TM_CCOEFF_NORMED)
_, maxv, _, maxl = cv2.minMaxLoc(res)
if maxv < TM_MIN_SCORE:
return None
bx = x1 + maxl[0]
by = y1 + maxl[1]
mcx = bx + tw / 2
mcy = by + th / 2
return (mcx, mcy, float(maxv), tw, th)
# =========================