diff --git a/ocv_track.py b/ocv_track.py index ed55417..c865c19 100644 --- a/ocv_track.py +++ b/ocv_track.py @@ -3,7 +3,7 @@ import sys import numpy as np import imutils as im -IMG_SCALE_UP = 2 +IMG_SCALE_UP = 4 TEMPLATE_SEARCH_AREA = 0.5 (major_ver, minor_ver, subminor_ver) = cv2.__version__.split('.') @@ -39,32 +39,58 @@ def process_image_bbox(src): return gray -def match(source, template, current_angle): +def match(source, template, angle, center): best_angle = 0 best_r = 0 - best_maxLoc = 0 + best_max_loc = 0 - image_source = im.resize(source, IMG_SCALE_UP * source.shape[0], IMG_SCALE_UP * source.shape[1]) + 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): - # scale up to IMG_SCALE_UP-size - image_template = im.resize(template, IMG_SCALE_UP * template.shape[0], IMG_SCALE_UP * template.shape[1]) - # rotated - image_template = im.rotate(image_template, current_angle + dangle) + # rotated and scale up to IMG_SCALE_UP-size + im_template = im.rotate(template, angle + dangle, center=center, scale=IMG_SCALE_UP) # Perform template match operations - res = cv2.matchTemplate(image_source, image_template, cv2.TM_CCOEFF_NORMED) + res = cv2.matchTemplate(im_source, im_template, cv2.TM_CCOEFF_NORMED) (minVal, maxVal, minLoc, maxLoc) = cv2.minMaxLoc(res) r = maxVal if r > best_r: best_r = r best_angle = dangle - best_maxLoc = maxLoc + best_max_loc = (maxLoc[0]/IMG_SCALE_UP, maxLoc[1]/IMG_SCALE_UP) - current_angle += best_angle + angle += best_angle - return current_angle, best_maxLoc + return angle, best_max_loc + + +class Reference: + def __init__(self, ref_gray): + self.center = (0, 0) + self.ref = cv2.cvtColor(ref_gray, cv2.COLOR_GRAY2BGR) + self.work = self.ref.copy() + cv2.namedWindow('Ref') + cv2.setMouseCallback('Ref', self.extract_coordinates) + + def set(self): + while True: + cv2.imshow("Ref", self.work) + key = cv2.waitKey(1) + if key == 32: + break + elif key == 27: + break + return self.center + + def extract_coordinates(self, event, x, y, flags, parameters): + if event == cv2.EVENT_LBUTTONDOWN: + self.work = self.ref.copy() + self.center = (x, y) + + print(f"Pixel at {self.center} = {self.work[y, x]}") + self.work[y, x] = (0, 0, 255) + print(f"Pixel at {self.center} = {self.work[y, x]}") trackers = { @@ -95,10 +121,12 @@ if not ok: print('Cannot read video file') sys.exit() -# Uncomment the line below to select a different bounding box bbox = cv2.selectROI("Frame", frame, False) image_template = process_image_bbox(image_crop(frame, bbox, TEMPLATE_SEARCH_AREA)) +ref = Reference(image_template) +ref_center = ref.set() +print(f"ref_center = {ref_center}") # Initialize tracker with first frame and bounding box tracker.create() tracker.init(frame, bbox) @@ -120,22 +148,17 @@ while True: image_template_rotated = image_template if ok: - current_angle, max_loc = match(image_bbox, image_template, current_angle) - image_template_rotated = im.rotate(image_template, current_angle) + 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 - startX += bbox[0] - startY += bbox[1] - endX = startX + bbox[2] - endY = startY + bbox[3] - print(f"startX: {startX}, startY: {startY}") + image_template_rotated = im.rotate(image_template, current_angle, center=ref_center) # draw the bounding box on the image - p1 = (int(bbox[0]), int(bbox[1])) - p2 = (int(bbox[0] + bbox[2]), int(bbox[1] + bbox[3])) + p1 = (int(startX + bbox[0]), int(startY + bbox[1])) + p2 = (int(bbox[0]+bbox[2]), int(bbox[1]+bbox[3])) cv2.rectangle(frame, p1, p2, (255, 0, 0), 3) # Calculate Frames per second (FPS)