Refactored

ocv_corner_tracker.py

@@ -57,7 +57,6 @@ class Corner:
 		matcher_bbox_local = (matcher_start_x, matcher_start_y, template.shape[1], template.shape[0])
 
 		cv2.imshow(f"{self.name}: Matcher res", matcher_res)
-#	cv2.imshow("{self.name}: Matcher view", matcher_crop)
 
 		return matcher_bbox_local, matcher_res
 
@@ -87,6 +86,22 @@ class CornerTracker:
 	def _print(self, s):
 		print(f"{self.name}: {s}")
 
+	def _debug(self, image_anno, matcher_bbox_local):
+		matcher_bbox = bbox_round(bbox_add_position(matcher_bbox_local, self.tracking_bb))
+		matcher_crop = image_crop(image, matcher_bbox)
+		cv2.imshow("{self.name}: Matcher view", matcher_crop)
+
+		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)
+		matcher_top_left_local, matcher_bottom_right_local = to_rect(bbox_round(matcher_bbox_local))
+
+		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)
+
 	def init_reference_frame(self, image: np.array):
 		print(f"Select tracking object")
 		self.tracking_ref_bb = cv2.selectROI("Image", image, False)
@@ -135,32 +150,28 @@ class CornerTracker:
 		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)
+			matcher_bbox = bbox_add_position(matcher_bbox_local, self.tracking_bb)
 
 			# 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)
+			self._debug(image_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]})")
@@ -181,7 +192,7 @@ class CornerTracker:
 		# 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]})")
+			self._print(f" -- Refined Corner  [{i}] ({corners_refined[i, 0, 0]}, {corners_refined[i, 0, 1]})")
 
 		return corners_refined
 
@@ -204,7 +215,8 @@ if __name__ == '__main__':
 			if ok:
 				ct.process(image)
 			else:
-				break
+				video.set(cv2.CAP_PROP_POS_FRAMES, 0)
+				continue
 
 			# Exit if ESC pressed
 			k = cv2.waitKey(key_wait) & 0xff