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.
61 lines
1.8 KiB
Python
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)
|
|
|
|
|
|
# =========================
|