import cv2 import numpy as np import imutils as im from util import bbox_extend, bbox_round, image_crop, to_rect, bbox_add_position, bbox_center 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 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 = [] self.line_from = None self.first_frame = True def _debug(self, image, matcher_res, matcher_start_x, matcher_start_y, template_scaled, crop_scaled): template_img = self.ref['image'] offset = self.ref['offset'] (min_val, max_val, min_loc, max_loc) = cv2.minMaxLoc(matcher_res) self._print(f"(min_val, max_val, min_loc, max_loc): {(min_val, max_val, min_loc, max_loc)}") self._print(f"matcher_start : {(matcher_start_x, matcher_start_y)}") self._print(f"matcher_res : {matcher_res.shape}") self._print(f"crop : {image.shape}") self._print(f"crop_scaled : {crop_scaled.shape}") self._print(f"template : {template_img.shape}") 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) def _print(self, s): print(f"{self.name}: {s}") def process(self, image: np.array): image_anno = image.copy() crop_scaled = im.resize(image, IMG_SCALE_UP * image.shape[1], IMG_SCALE_UP * image.shape[0]) template = self.ref['image'] bb_template = self.ref['bbox'] # Reseize template template_scaled = im.resize(template, IMG_SCALE_UP * template.shape[0], IMG_SCALE_UP * template.shape[1]) # Match matcher_res = cv2.matchTemplate(crop_scaled, template_scaled, cv2.TM_CCOEFF_NORMED) (min_val, max_val, min_loc, max_loc) = cv2.minMaxLoc(matcher_res) # Get matcher coord matcher_start_x = max_loc[0]/IMG_SCALE_UP matcher_start_y = max_loc[1]/IMG_SCALE_UP # Store offset at first frame if self.first_frame: self.first_frame = False offset_x = matcher_start_x - bb_template[0] offset_y = matcher_start_y - bb_template[1] self.ref['offset'] = (offset_x, offset_y) # Compensate offset matcher_start_x -= self.ref['offset'][0] matcher_start_y -= self.ref['offset'][1] matcher_bbox_local = (matcher_start_x, matcher_start_y, template.shape[1], template.shape[0]) self._debug(image_anno, matcher_res, matcher_start_x, matcher_start_y, template_scaled, crop_scaled) return matcher_bbox_local, matcher_res def path_add(self, point): if self.line_from is not None: self._path.append({'from': self.line_from, 'to': point}) self.line_from = point @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 _debug(self, image_anno, tracking_anno, matcher_bbox_local): matcher_bbox = bbox_round(bbox_add_position(matcher_bbox_local, self.tracking_bb)) template_top_left, template_bottom_right = to_rect(self.matching_tpl_bb) cv2.rectangle(tracking_anno, template_top_left, template_bottom_right, COLOR_TEMPLATE, 1) matcher_top_left_local, matcher_bottom_right_local = to_rect(bbox_round(matcher_bbox_local)) cv2.rectangle(tracking_anno, matcher_top_left_local, matcher_bottom_right_local, COLOR_MATCHER, 1) matcher_rect = to_rect(matcher_bbox) cv2.rectangle(image_anno, matcher_rect[0], matcher_rect[1], COLOR_MATCHER, 1) tracker_rect = to_rect(self.tracking_bb) cv2.rectangle(image_anno, tracker_rect[0], tracker_rect[1], COLOR_TRACKER, 1) def init_reference_frame(self, image: np.array): print(f"Select tracking object") self.tracking_ref_bb = cv2.selectROI("Select", 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, name=f"Corner-{count}")) 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) cv2.destroyWindow("Select") 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() tracking_anno = self.tracking_img.copy() corners = [] for ct in self.corner_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]) 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)) # store refined corners for i in range(corners_refined.shape[0]): line_to = (corners_refined[i, 0, 0], corners_refined[i, 0, 1]) ct = self.corner_list[i] ct.path_add(line_to) # draw path for ct in self.corner_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)) 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 Corner [{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: video.set(cv2.CAP_PROP_POS_FRAMES, 0) continue # 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')