From 8c499903dfc692aa798b6b1b1839702fc43c74d8 Mon Sep 17 00:00:00 2001
From: Jens Ahrensfeld <jens@jayfield.org>
Date: Thu, 4 Jul 2024 15:59:54 +0200
Subject: [PATCH] added multi tracking

---
 ocv_corner_tracker.py | 47 +++++++++++++++++++++++++++----------------
 1 file changed, 30 insertions(+), 17 deletions(-)

diff --git a/ocv_corner_tracker.py b/ocv_corner_tracker.py
index 3d3bf44..8058f7b 100644
--- a/ocv_corner_tracker.py
+++ b/ocv_corner_tracker.py
@@ -91,7 +91,8 @@ class Corner:
 
 
 class CornerTracker:
-	def __init__(self, name: str = 'CornerTracker'):
+	def __init__(self, color=(0, 255, 0), name: str = 'CornerTracker'):
+		self.color = color
 		self.name = name
 		self.tracking_ref_bb = None
 		self.tracking_ref_img = None
@@ -121,7 +122,7 @@ class CornerTracker:
 		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)
+		cv2.rectangle(image_anno, tracker_rect[0], tracker_rect[1], self.color, 1)
 
 	def init_reference_frame(self, image: np.array):
 		print(f"Select tracking object")
@@ -155,9 +156,9 @@ class CornerTracker:
 
 		# Refine initial corners and store them as reference
 		self.corner_ref = self._corner_refine(cv2.cvtColor(image, cv2.COLOR_BGR2GRAY), corners=corner_list)
-		cv2.destroyWindow("Select")
+		#cv2.destroyWindow("Select")
 
-	def process(self, image: np.array):
+	def process(self, image: np.array, image_anno: np.array):
 		if self.tracker is None:
 			raise Exception(f"{self.name}: Call init_reference_frame() first")
 
@@ -171,10 +172,9 @@ class CornerTracker:
 		self.tracking_img = cv2.GaussianBlur(tracking_img, (9, 9), 0)
 		self.tracking_bb = tracking_bb
 
-		return self._match(image)
+		return self._match(image, image_anno)
 
-	def _match(self, image: np.array):
-		image_anno = image.copy()
+	def _match(self, image: np.array, image_anno: np.array):
 		tracking_anno = self.tracking_img.copy()
 		corners_raw = []
 		for ct in self.corner_matcher_list:
@@ -209,11 +209,12 @@ class CornerTracker:
 			self._print(f" -- Corner fine      [{i}] {corners_refined[i]}")
 			self._print(f" -- Corner distance  [{i}] {corners_refined[i] - self.corner_ref[i]}")
 			cv2.circle(image_anno, (int(corners_refined[i, 0]), int(corners_refined[i, 1])), 4, (0, 255, 0))
+			cv2.putText(image_anno, f"Distance [{i}] : {corners_refined[i] - self.corner_ref[i]}", (50, 20 + 20*i), cv2.FONT_HERSHEY_SIMPLEX, 0.5, self.color, 1)
 
 		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: np.array):
+	@staticmethod
+	def _corner_refine(src_gray, corners: np.array):
 		# convert from (n, 2) to (n, 1, 2)
 		corners_original = []
 		for corner in corners:
@@ -242,22 +243,34 @@ if __name__ == '__main__':
 	float_formatter = "{:.3f}".format
 	np.set_printoptions(formatter={'float_kind': float_formatter})
 
+	colors = [(255, 0, 0), (0, 255, 0), (0, 0, 255), (255, 0, 255), (0, 255, 255), (255, 255, 255)]
 	video = cv2.VideoCapture('./data/spindle_multi_black/%04d.png')
+	tracker_count = 0
+	tracker_list = []
 	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')
+		for tracker_count in range(0, 6):
+			ct = CornerTracker(colors[tracker_count], name=f"Tracker-{tracker_count}")
+			# Read first frame.
+			ok, image = video.read()
+			if ok:
+				ct.init_reference_frame(image)
+			else:
+				print('Cannot read video file')
+
+			tracker_list.append(ct)
+			k = cv2.waitKey(-1) & 0xff
+			if k == 27:
+				break
 
 		key_wait = -1
 		while True:
 			# Read a new frame
 			ok, image = video.read()
 			if ok:
-				ct.process(image)
+				image_anno = image.copy()
+				for ct in tracker_list:
+					ct.process(image, image_anno)
+				cv2.imshow(f"Image Anno", image_anno)
 			else:
 				video.set(cv2.CAP_PROP_POS_FRAMES, 0)
 				continue