diff --git a/ocv_corner_tracker.py b/ocv_corner_tracker.py index 4164c31..a19b553 100644 --- a/ocv_corner_tracker.py +++ b/ocv_corner_tracker.py @@ -7,9 +7,15 @@ from util import bbox_extend, bbox_round, image_crop, to_rect, bbox_add_position IMG_SCALE_UP = 1 IMG_ROTATE = 0 +COLOR_TRACKER = (0, 255, 0) +COLOR_TRACKER_EXT = (0, 255, 255) +COLOR_MATCHER = (0, 0, 255) +COLOR_TEMPLATE = (255, 0, 0) +TEMPLATE_MATCH_OVERLAP = 0 -class CornerTracker: - def __init__(self, reference: np.array, bbox: np.array, name: str = 'CornerTracker'): + +class Corner: + def __init__(self, reference: np.array, bbox: np.array, name: str = 'Corner'): self.ref = {'image': reference, 'bbox': bbox, 'offset': (0, 0)} self.name = name self._path = [] @@ -63,3 +69,152 @@ class CornerTracker: @property def path(self): return self._path + + +class CornerTracker: + def __init__(self, name: str = 'CornerTracker'): + 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 = None + self.corner_list = [] + + def _print(self, s): + print(f"{self.name}: {s}") + + def init_reference_frame(self, image: np.array): + print(f"Select tracking object") + self.tracking_ref_bb = cv2.selectROI("Image", image, 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) + + self.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)) + self._print(f"Corner {count} added") + self._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") + + # Initialize tracker with first frame and bounding box + self.tracker = cv2.TrackerKCF().create() + self.tracker.init(image, self.tracking_ref_bb) + + def process(self, image: np.array): + if self.tracker is None: + raise Exception(f"{self.name}: Call init_reference_frame() first") + + # Update tracker + _ok, tracking_bb = self.tracker.update(image) + if not _ok: + return None + + tracking_img = image_crop(image.copy(), 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) + + def _match(self, image: np.array): + image_anno = image.copy() + corners = [] + for ct in self.corner_list: + matcher_bbox_local, matcher = ct.process(self.tracking_img) + matcher_top_left_local, matcher_bottom_right_local = to_rect(bbox_round(matcher_bbox_local)) + + matcher_bbox = bbox_round(bbox_add_position(matcher_bbox_local, self.tracking_bb)) + matcher_rect = to_rect(matcher_bbox) + + tracker_rect = to_rect(self.tracking_bb) + template_top_left, template_bottom_right = to_rect(self.matching_tpl_bb) + + cv2.rectangle(self.tracking_img, template_top_left, template_bottom_right, COLOR_TEMPLATE, 1) + cv2.rectangle(self.tracking_img, matcher_top_left_local, matcher_bottom_right_local, COLOR_MATCHER, 1) + + cv2.rectangle(image_anno, matcher_rect[0], matcher_rect[1], COLOR_MATCHER, 1) + cv2.rectangle(image_anno, tracker_rect[0], tracker_rect[1], COLOR_TRACKER, 1) + matcher_crop = image_crop(image, matcher_bbox) + + # Draw path + line_to = bbox_center(matcher_bbox) + corners.append([line_to]) + ct.path_add(line_to) + + for p in ct.path: + cv2.line(image_anno, bbox_round(p['from']), bbox_round(p['to']), (0, 0, 255), 2) + + # refine corners + corners_refined = self._corner_refine(cv2.cvtColor(image, cv2.COLOR_BGR2GRAY), np.array(corners, dtype=np.float32)) + + # 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)) + + 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): + # Set the needed parameters to find the refined corners + win_size = (5, 5) + zero_zone = (-1, -1) + criteria = (cv2.TERM_CRITERIA_EPS + cv2.TermCriteria_COUNT, 40, 0.001) + + # Calculate the refined corner locations + corners_refined = cv2.cornerSubPix(src_gray, corners_original.copy(), win_size, zero_zone, criteria) + + # Write them down + 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 Corne r [{i}] ({corners_refined[i, 0, 0]}, {corners_refined[i, 0, 1]})") + + return corners_refined + + +if __name__ == '__main__': + video = cv2.VideoCapture('./data/spindle_multi_black/%04d.png') + if video.isOpened(): + ct = CornerTracker() + # Read first frame. + ok, image = video.read() + if ok: + ct.init_reference_frame(image) + else: + print('Cannot read video file') + + key_wait = -1 + while True: + # Read a new frame + ok, image = video.read() + if ok: + ct.process(image) + else: + break + + # Exit if ESC pressed + k = cv2.waitKey(key_wait) & 0xff + if k == 27: + break + if k == ord(' '): + if key_wait == -1: + key_wait = 30 + else: + key_wait = -1 + + else: + print('Cannot open video file')