diff --git a/ocv_corner_tracker.py b/ocv_corner_tracker.py index 0472bc2..5fbbc18 100644 --- a/ocv_corner_tracker.py +++ b/ocv_corner_tracker.py @@ -18,7 +18,7 @@ COLOR_CIRCLE = (0, 255, 0) TEMPLATE_MATCH_OVERLAP = 0 CONSOLE_DEBUG = False -IMAGE_DEBUG = False +IMAGE_DEBUG = True DO_TRACKING = False @@ -105,9 +105,7 @@ class CornerTracker: self.color = color self.name = name self.matching_tpl_bb = None - self.matching_tpl_img = None self.tracking_bb = None - self.tracking_img = None self.tracking_mask = None self.tracker = None self.corner_ref = None @@ -157,6 +155,13 @@ class CornerTracker: _image[_mask == 0] = 0 _image[_mask != 0] = _image[_mask != 0] + @staticmethod + def image_process(_image: np.array): + result = cv2.cvtColor(_image, cv2.COLOR_BGR2GRAY) + result = cv2.GaussianBlur(result, (9, 9), 0) + + return result + def init_reference_frame(self, _image: np.array): # Draw mask _mask = CornerTracker.mask_init(image) @@ -169,27 +174,29 @@ class CornerTracker: if bbox_center(bb) == (0, 0): return False - self.tracking_bb = bb - self.tracking_img = image_crop(_image, self.tracking_bb) - self.tracking_mask = image_crop(_mask, self.tracking_bb) - tracking_ref_gray_img = cv2.cvtColor(self.tracking_img, cv2.COLOR_BGR2GRAY) - self.corner_matcher_list = [] + tracking_img = image_crop(_image, bb) + image_process_local = CornerTracker.image_process(tracking_img) + tracking_mask = image_crop(_mask, bb) + masked_tracking_img = tracking_img.copy() + CornerTracker.mask_apply(masked_tracking_img, tracking_mask) - masked_tracking_img = self.tracking_img.copy() - CornerTracker.mask_apply(masked_tracking_img, self.tracking_mask) + self.corner_matcher_list = [] + self.tracking_bb = bb + self.tracking_mask = tracking_mask corner_list = [] count = 1 while True: print(f"Add Corner {count}") - bb = cv2.selectROI("Matcher Reference", masked_tracking_img, fromCenter=True, showCrosshair=True) - if bbox_center(bb) == (0, 0): + matching_tpl_bb = cv2.selectROI("Matcher Reference", masked_tracking_img, fromCenter=True, showCrosshair=True) + if bbox_center(matching_tpl_bb) == (0, 0): break - self.matching_tpl_bb = 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_bb))) + matching_tpl_img = image_crop(image_process_local.copy(), matching_tpl_bb) + self.corner_matcher_list.append(Corner(matching_tpl_img, matching_tpl_bb, name=f"Corner-{count}")) + corner_list.append(bbox_center(matching_tpl_bb)) + + self.matching_tpl_bb = matching_tpl_bb print(f"Corner {count} added") print(f"Press any key to add another corner or ESC to continue") @@ -207,8 +214,17 @@ class CornerTracker: self.tracker = cv2.TrackerKCF.create() self.tracker.init(_image, self.tracking_bb) + print(corner_list) + cv2.imshow("image_process_local", image_process_local) + # Refine initial corners and store them as reference - self.corner_ref = self._corner_refine(cv2.cvtColor(_image, cv2.COLOR_BGR2GRAY), corners=corner_list) + corners_local = self._corner_refine(image_process_local, corners=corner_list) + + # Transform corners to global + self.corner_ref = [] + for corner_local in corners_local: + corner = (corner_local[0] + self.tracking_bb[0], corner_local[1] + self.tracking_bb[1]) + self.corner_ref.append(corner) return True @@ -221,17 +237,14 @@ class CornerTracker: _image_local = image_crop(_image.copy(), self.tracking_bb).copy() CornerTracker.mask_apply(_image_local, self.tracking_mask) - tracking_img = cv2.cvtColor(_image_local, cv2.COLOR_BGR2GRAY) - CornerTracker.mask_apply(tracking_img, self.tracking_mask) - self.tracking_img = cv2.GaussianBlur(tracking_img, (9, 9), 0) + _image_processed_local = CornerTracker.image_process(_image_local) - return self._match(_image_local, _image_anno) + return self._match(_image_processed_local, _image_anno, _image_local.copy()) - def _match(self, _image: np.array, _image_anno: np.array): - tracking_anno = self.tracking_img.copy() + def _match(self, _image: np.array, _image_anno: np.array, tracking_anno): corners_raw = [] for _ct in self.corner_matcher_list: - matcher_bbox_local, matcher = _ct.process(self.tracking_img) + matcher_bbox_local, matcher = _ct.process(_image) # Draw path corner = bbox_center(matcher_bbox_local) @@ -240,11 +253,11 @@ class CornerTracker: self._debug(_image_anno, tracking_anno, matcher_bbox_local) # refine corners - corners_refined_local = self._corner_refine(cv2.cvtColor(_image, cv2.COLOR_BGR2GRAY), corners_raw) + corners_local = self._corner_refine(_image, corners_raw) # Transform corners to global corners_refined = [] - for corner_local in corners_refined_local: + for corner_local in corners_local: corner = (corner_local[0] + self.tracking_bb[0], corner_local[1] + self.tracking_bb[1]) corners_refined.append(corner) @@ -277,7 +290,7 @@ class CornerTracker: cv2.circle(_image_anno, (int(corners_refined[_i, 0] + 0.5), int(corners_refined[_i, 1] + 0.5)), 4, COLOR_CIRCLE) if IMAGE_DEBUG: - cv2.imshow(f"{self.name}: Tracker: ", self.tracking_img) + cv2.imshow(f"{self.name}: Tracker: ", tracking_anno) return _mean_distance @@ -342,6 +355,9 @@ if __name__ == '__main__': else: print('Cannot read video file') + # rewind to frame 0 + video.set(cv2.CAP_PROP_POS_FRAMES, 0) + key_wait = -1 dist_min = [(0, 0)]*len(tracker_list) dist_max = [(0, 0)]*len(tracker_list)