diff --git a/ocv_corner_tracker.py b/ocv_corner_tracker.py index 8058f7b..343a901 100644 --- a/ocv_corner_tracker.py +++ b/ocv_corner_tracker.py @@ -14,6 +14,8 @@ COLOR_TEMPLATE = (255, 0, 0) TEMPLATE_MATCH_OVERLAP = 0 CONSOLE_DEBUG = False +IMAGE_DEBUG = False + class Corner: def __init__(self, reference: np.array, bbox: np.array, name: str = 'Corner'): @@ -38,10 +40,11 @@ class Corner: self._print(f"template_scaled : {template_scaled.shape}") self._print(f"offset : {offset}") - matcher_bbox_local = (matcher_start_x, matcher_start_y, template_img.shape[1], template_img.shape[0]) - matcher_crop = image_crop(image, matcher_bbox_local) - cv2.imshow(f"{self.name}: Matcher res", matcher_res) - cv2.imshow(f"{self.name}: Matcher view", matcher_crop) + if IMAGE_DEBUG: + matcher_bbox_local = (matcher_start_x, matcher_start_y, template_img.shape[1], template_img.shape[0]) + matcher_crop = image_crop(image, matcher_bbox_local) + cv2.imshow(f"{self.name}: Matcher res", matcher_res) + cv2.imshow(f"{self.name}: Matcher view", matcher_crop) def _print(self, s): if CONSOLE_DEBUG: @@ -124,9 +127,9 @@ class CornerTracker: tracker_rect = to_rect(self.tracking_bb) cv2.rectangle(image_anno, tracker_rect[0], tracker_rect[1], self.color, 1) - def init_reference_frame(self, image: np.array): + def init_reference_frame(self, image: np.array, select_window: np.array): print(f"Select tracking object") - self.tracking_ref_bb = cv2.selectROI("Select", image, False) + self.tracking_ref_bb = cv2.selectROI("Tracker Reference", select_window, False) self.tracking_ref_img = image_crop(image.copy(), self.tracking_ref_bb) self.tracking_ref_gray_img = cv2.cvtColor(self.tracking_ref_img, cv2.COLOR_BGR2GRAY) @@ -134,21 +137,21 @@ class CornerTracker: corner_list = [] count = 1 while True: - self._print(f"Add Corner {count}") - self.matching_tpl_bb = cv2.selectROI("Image", self.tracking_ref_img, fromCenter=True, showCrosshair=True) + print(f"Add Corner {count}") + self.matching_tpl_bb = cv2.selectROI("Matcher Reference", self.tracking_ref_img, fromCenter=True, showCrosshair=True) self.matching_tpl_img = image_crop(self.tracking_ref_gray_img.copy(), self.matching_tpl_bb) self.corner_matcher_list.append(Corner(self.matching_tpl_img, self.matching_tpl_bb, name=f"Corner-{count}")) corner_list.append(bbox_center(bbox_add_position(self.matching_tpl_bb, self.tracking_ref_bb))) - self._print(f"Corner {count} added") - self._print(f"Press any key to add another corner or ESC to continue") + print(f"Corner {count} added") + print(f"Press any key to add another corner or ESC to continue") k = cv2.waitKey(-1) & 0xff if k == 27: break count += 1 - self._print(f"Added {count} corners") + print(f"Added {count} corners") # Initialize tracker with first frame and bounding box self.tracker = cv2.TrackerKCF().create() @@ -156,7 +159,6 @@ class CornerTracker: # Refine initial corners and store them as reference self.corner_ref = self._corner_refine(cv2.cvtColor(image, cv2.COLOR_BGR2GRAY), corners=corner_list) - #cv2.destroyWindow("Select") def process(self, image: np.array, image_anno: np.array): if self.tracker is None: @@ -202,6 +204,12 @@ class CornerTracker: for p in ct.path: cv2.line(image_anno, bbox_round(p['from']), bbox_round(p['to']), (0, 0, 255), 1) + distances = [] + for i in range(0, corners_refined.shape[0]): + distances.append((corners_refined[i] - self.corner_ref[i])) + + mean_distance = np.mean(distances, axis=0) + # show refined corners for i in range(0, corners_refined.shape[0]): self._print(f" -- Corner Reference [{i}] {self.corner_ref[i]}") @@ -209,9 +217,11 @@ class CornerTracker: self._print(f" -- Corner fine [{i}] {corners_refined[i]}") self._print(f" -- Corner distance [{i}] {corners_refined[i] - self.corner_ref[i]}") cv2.circle(image_anno, (int(corners_refined[i, 0]), int(corners_refined[i, 1])), 4, (0, 255, 0)) - cv2.putText(image_anno, f"Distance [{i}] : {corners_refined[i] - self.corner_ref[i]}", (50, 20 + 20*i), cv2.FONT_HERSHEY_SIMPLEX, 0.5, self.color, 1) - cv2.imshow(f"{self.name}: Tracker: ", self.tracking_img) + if IMAGE_DEBUG: + cv2.imshow(f"{self.name}: Tracker: ", self.tracking_img) + + return mean_distance @staticmethod def _corner_refine(src_gray, corners: np.array): @@ -248,19 +258,23 @@ if __name__ == '__main__': tracker_count = 0 tracker_list = [] if video.isOpened(): - for tracker_count in range(0, 6): - ct = CornerTracker(colors[tracker_count], name=f"Tracker-{tracker_count}") - # Read first frame. - ok, image = video.read() - if ok: - ct.init_reference_frame(image) - else: - print('Cannot read video file') + # Read first frame. + ok, image = video.read() + select_window = image.copy() + if ok: + for tracker_count in range(0, 6): + print(f"Add tracker #{tracker_count}") + ct = CornerTracker(colors[tracker_count], name=f"Tracker-{tracker_count}") + ct.init_reference_frame(image, select_window) + tracker_list.append(ct) + print(f"Press any key to add another tracker or ESC to continue") + k = cv2.waitKey(-1) & 0xff + if k == 27: + break - tracker_list.append(ct) - k = cv2.waitKey(-1) & 0xff - if k == 27: - break + print(f"Number of active tracker: {len(tracker_list)}") + else: + print('Cannot read video file') key_wait = -1 while True: @@ -268,8 +282,12 @@ if __name__ == '__main__': ok, image = video.read() if ok: image_anno = image.copy() + i = 0 for ct in tracker_list: - ct.process(image, image_anno) + mean_distance = ct.process(image, image_anno) + cv2.putText(image_anno, f"Distance [{i}] : {mean_distance}", (50, 20 + 20 * i), + cv2.FONT_HERSHEY_SIMPLEX, 0.6, ct.color, 1) + i += 1 cv2.imshow(f"Image Anno", image_anno) else: video.set(cv2.CAP_PROP_POS_FRAMES, 0)