diff --git a/ocv_corner_tracker.py b/ocv_corner_tracker.py index b3c0f4e..3d3bf44 100644 --- a/ocv_corner_tracker.py +++ b/ocv_corner_tracker.py @@ -13,6 +13,7 @@ COLOR_MATCHER = (0, 0, 255) COLOR_TEMPLATE = (255, 0, 0) TEMPLATE_MATCH_OVERLAP = 0 +CONSOLE_DEBUG = False class Corner: def __init__(self, reference: np.array, bbox: np.array, name: str = 'Corner'): @@ -43,7 +44,8 @@ class Corner: cv2.imshow(f"{self.name}: Matcher view", matcher_crop) def _print(self, s): - print(f"{self.name}: {s}") + if CONSOLE_DEBUG: + print(f"{self.name}: {s}") def process(self, image: np.array): image_anno = image.copy() @@ -99,10 +101,12 @@ class CornerTracker: self.tracking_bb = None self.tracking_img = None self.tracker = None - self.corner_list = [] + self.corner_ref = None + self.corner_matcher_list = [] def _print(self, s): - print(f"{self.name}: {s}") + if CONSOLE_DEBUG: + print(f"{self.name}: {s}") def _debug(self, image_anno, tracking_anno, matcher_bbox_local): matcher_bbox = bbox_round(bbox_add_position(matcher_bbox_local, self.tracking_bb)) @@ -125,13 +129,16 @@ class CornerTracker: 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) - self.corner_list = [] + self.corner_matcher_list = [] + 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) self.matching_tpl_img = image_crop(self.tracking_ref_gray_img.copy(), self.matching_tpl_bb) - self.corner_list.append(Corner(self.matching_tpl_img, self.matching_tpl_bb, name=f"Corner-{count}")) + 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") @@ -145,6 +152,9 @@ class CornerTracker: # Initialize tracker with first frame and bounding box self.tracker = cv2.TrackerKCF().create() self.tracker.init(image, self.tracking_ref_bb) + + # 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): @@ -166,40 +176,51 @@ class CornerTracker: def _match(self, image: np.array): image_anno = image.copy() tracking_anno = self.tracking_img.copy() - corners = [] - for ct in self.corner_list: + corners_raw = [] + for ct in self.corner_matcher_list: matcher_bbox_local, matcher = ct.process(self.tracking_img) matcher_bbox = bbox_add_position(matcher_bbox_local, self.tracking_bb) # Draw path - line_to = bbox_center(matcher_bbox) - corners.append([line_to]) + corner = bbox_center(matcher_bbox) + corners_raw.append(corner) self._debug(image_anno, tracking_anno, matcher_bbox_local) # refine corners - corners_refined = self._corner_refine(cv2.cvtColor(image, cv2.COLOR_BGR2GRAY), np.array(corners, dtype=np.float32)) + corners_refined = self._corner_refine(cv2.cvtColor(image, cv2.COLOR_BGR2GRAY), corners_raw) # store refined corners - for i in range(corners_refined.shape[0]): - line_to = (int(corners_refined[i, 0, 0]), int(corners_refined[i, 0, 1])) - ct = self.corner_list[i] - ct.path_add(line_to) + i = 0 + for corner in corners_refined: + ct = self.corner_matcher_list[i] + ct.path_add(corner) + i += 1 # draw path - for ct in self.corner_list: + for ct in self.corner_matcher_list: for p in ct.path: cv2.line(image_anno, bbox_round(p['from']), bbox_round(p['to']), (0, 0, 255), 1) # show refined corners - for i in range(corners_refined.shape[0]): - self._print(f" -- Refined Corner [{i}] ({corners_refined[i, 0, 0]}, {corners_refined[i, 0, 1]})") - cv2.circle(image_anno, (int(corners_refined[i, 0, 0]), int(corners_refined[i, 0, 1])), 4, (0, 255, 0)) + for i in range(0, corners_refined.shape[0]): + self._print(f" -- Corner Reference [{i}] {self.corner_ref[i]}") + self._print(f" -- Corner coarse [{i}] {corners_raw[i]}") + 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.imshow(f"{self.name}: Tracker: ", self.tracking_img) cv2.imshow(f"{self.name}: Image Anno", image_anno) - def _corner_refine(self, src_gray, corners_original: np.array): + def _corner_refine(self, src_gray, corners: np.array): + # convert from (n, 2) to (n, 1, 2) + corners_original = [] + for corner in corners: + corners_original.append([corner]) + + corners_original = np.array(corners_original, dtype=np.float32) + # Set the needed parameters to find the refined corners win_size = (5, 5) zero_zone = (-1, -1) @@ -209,14 +230,18 @@ class CornerTracker: corners_refined = cv2.cornerSubPix(src_gray, corners_original.copy(), win_size, zero_zone, criteria) # Write them down + corner_result = [] for i in range(corners_original.shape[0]): - self._print(f" -- Original Corner [{i}] ({corners_original[i, 0, 0]}, {corners_original[i, 0, 1]})") - self._print(f" -- Refined Corner [{i}] ({corners_refined[i, 0, 0]}, {corners_refined[i, 0, 1]})") + corner = (corners_refined[i, 0, 0], corners_refined[i, 0, 1]) + corner_result.append(corner) - return corners_refined + return np.array(corner_result) if __name__ == '__main__': + float_formatter = "{:.3f}".format + np.set_printoptions(formatter={'float_kind': float_formatter}) + video = cv2.VideoCapture('./data/spindle_multi_black/%04d.png') if video.isOpened(): ct = CornerTracker()