diff --git a/ocv_track.py b/ocv_track.py index c865c19..94ebd3e 100644 --- a/ocv_track.py +++ b/ocv_track.py @@ -4,7 +4,7 @@ import numpy as np import imutils as im IMG_SCALE_UP = 4 -TEMPLATE_SEARCH_AREA = 0.5 +TEMPLATE_SEARCH_AREA = 0.8 (major_ver, minor_ver, subminor_ver) = cv2.__version__.split('.') print(cv2.__version__) @@ -21,8 +21,19 @@ def bbox_extend(bbox: cv2.typing.Rect, search_area: float): return res +def bbox_round(src): + x = int(round(src[0])) + y = int(round(src[1])) + w = int(round(src[2])) + h = int(round(src[3])) + + return x, y, w, h + + def image_crop(src, bbox: cv2.typing.Rect, search_area: float = 0): - bbox = bbox_extend(bbox, search_area) + if search_area > 0: + bbox = bbox_extend(bbox, search_area) + x = bbox[0] y = bbox[1] w = bbox[2] @@ -39,16 +50,15 @@ def process_image_bbox(src): return gray -def match(source, template, angle, center): +def template_match_rotation(source, template, angle, center, rot_min=-1.0, rot_max=+1.0, n_steps=10, scale=1): best_angle = 0 best_r = 0 best_max_loc = 0 - - im_source = im.resize(source, IMG_SCALE_UP * source.shape[0], IMG_SCALE_UP * source.shape[1]) - - for dangle in np.linspace(-1.0, +1.0, 10): + r = 0 + im_source = im.resize(source, scale * source.shape[0], scale * source.shape[1]) + for dangle in np.linspace(rot_min, rot_max, n_steps): # rotated and scale up to IMG_SCALE_UP-size - im_template = im.rotate(template, angle + dangle, center=center, scale=IMG_SCALE_UP) + im_template = im.rotate(template, angle + dangle, center=center, scale=scale) # Perform template match operations res = cv2.matchTemplate(im_source, im_template, cv2.TM_CCOEFF_NORMED) @@ -58,13 +68,24 @@ def match(source, template, angle, center): if r > best_r: best_r = r best_angle = dangle - best_max_loc = (maxLoc[0]/IMG_SCALE_UP, maxLoc[1]/IMG_SCALE_UP) + best_max_loc = (maxLoc[0]/scale, maxLoc[1]/scale) + else: + break angle += best_angle return angle, best_max_loc +def match(source, template, angle, center): + # coarse + angle_c, loc = template_match_rotation(source, template, angle, center) + # fine + angle, loc = template_match_rotation(source, template, angle_c, center, rot_min=-0.25, rot_max=+0.25, n_steps=20, scale=IMG_SCALE_UP) + + return angle, loc + + class Reference: def __init__(self, ref_gray): self.center = (0, 0) @@ -104,7 +125,7 @@ trackers = { 'BOOSTING': cv2.legacy.TrackerBoosting, 'MEDIANFLOW': cv2.legacy.TrackerMedianFlow } -tracker_type = 'MOSSE' +tracker_type = 'KCF' tracker = trackers[tracker_type].create() video = cv2.VideoCapture("./data/production_id 4525346 (1080p).mp4") @@ -142,18 +163,19 @@ while True: timer = cv2.getTickCount() # Update tracker - ok, bbox = tracker.update(frame) + ok, bbox_tracker = tracker.update(frame) - image_bbox = process_image_bbox(image_crop(frame, bbox)) image_template_rotated = image_template + bbox = bbox_round(bbox_tracker) if ok: + image_bbox = process_image_bbox(image_crop(frame, bbox)) current_angle, max_loc = match(image_bbox, image_template, current_angle, ref_center) # determine the starting and ending (x, y)-coordinates of the # bounding box (startX, startY) = max_loc - print(f"startX: {startX}, startY: {startY}") + print(f"current_angle: {current_angle:.2f}, startX: {startX}, startY: {startY}") image_template_rotated = im.rotate(image_template, current_angle, center=ref_center) # draw the bounding box on the image