CornerTracker object oriented

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')