diff --git a/ocv_corner_tracker.py b/ocv_corner_tracker.py index fe9d0ad..68a6257 100644 --- a/ocv_corner_tracker.py +++ b/ocv_corner_tracker.py @@ -15,6 +15,7 @@ TEMPLATE_MATCH_OVERLAP = 0 CONSOLE_DEBUG = False IMAGE_DEBUG = False +DO_TRACKING = False class Corner: @@ -99,14 +100,11 @@ class CornerTracker: def __init__(self, color=(0, 255, 0), name: str = 'CornerTracker'): self.color = color self.name = name - self.tracking_ref_bb = None - self.tracking_ref_img = None - self.tracking_ref_gray_img = None self.matching_tpl_bb = None self.matching_tpl_img = None self.tracking_bb = None self.tracking_img = None - self.tracker = cv2.TrackerKCF.create() + self.tracker = None self.corner_ref = None self.corner_matcher_list = [] @@ -136,23 +134,23 @@ class CornerTracker: if bbox_center(bb) == (0, 0): return False - self.tracking_ref_bb = bb - self.tracking_ref_img = image_crop(_image, self.tracking_ref_bb) - self.tracking_ref_gray_img = cv2.cvtColor(self.tracking_ref_img, cv2.COLOR_BGR2GRAY) + self.tracking_bb = bb + self.tracking_img = image_crop(_image, self.tracking_bb) + tracking_ref_gray_img = cv2.cvtColor(self.tracking_img, cv2.COLOR_BGR2GRAY) self.corner_matcher_list = [] corner_list = [] count = 1 while True: print(f"Add Corner {count}") - bb = cv2.selectROI("Matcher Reference", self.tracking_ref_img, fromCenter=True, showCrosshair=True) + bb = cv2.selectROI("Matcher Reference", self.tracking_img, fromCenter=True, showCrosshair=True) if bbox_center(bb) == (0, 0): break self.matching_tpl_bb = bb - self.matching_tpl_img = image_crop(self.tracking_ref_gray_img.copy(), self.matching_tpl_bb) + self.matching_tpl_img = image_crop(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))) + corner_list.append(bbox_center(bbox_add_position(self.matching_tpl_bb, self.tracking_bb))) print(f"Corner {count} added") print(f"Press any key to add another corner or ESC to continue") @@ -166,7 +164,9 @@ class CornerTracker: return False # Initialize tracker with first frame and bounding box - self.tracker.init(_image, self.tracking_ref_bb) + if DO_TRACKING: + self.tracker = cv2.TrackerKCF.create() + self.tracker.init(_image, self.tracking_bb) # Refine initial corners and store them as reference self.corner_ref = self._corner_refine(cv2.cvtColor(_image, cv2.COLOR_BGR2GRAY), corners=corner_list) @@ -174,18 +174,15 @@ class CornerTracker: return True def process(self, _image: np.array, _image_anno: np.array): - if self.tracker is None: - raise Exception(f"{self.name}: Call init_reference_frame() first") + # Update tracker bb + if self.tracker is not None: + _ok, self.tracking_bb = self.tracker.update(_image) + if not _ok: + return None - # Update tracker - _ok, tracking_bb = self.tracker.update(_image) - if not _ok: - return None - - tracking_img = image_crop(_image.copy(), tracking_bb) + tracking_img = image_crop(_image.copy(), self.tracking_bb) tracking_img = cv2.cvtColor(tracking_img, cv2.COLOR_BGR2GRAY) self.tracking_img = cv2.GaussianBlur(tracking_img, (9, 9), 0) - self.tracking_bb = tracking_bb return self._match(_image, _image_anno)