Refactored

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)