Reformatted file

ocv_corner_tracker.py

@@ -2,7 +2,7 @@ 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
+from util import bbox_round, image_crop, to_rect, bbox_add_position, bbox_center
 
 IMG_SCALE_UP = 1
 IMG_ROTATE = 0
@@ -25,7 +25,7 @@ class Corner:
 		self.line_from = None
 		self.first_frame = True
 
-	def _debug(self, image, matcher_res, matcher_start_x, matcher_start_y, template_scaled, crop_scaled):
+	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']
 
@@ -34,7 +34,7 @@ class Corner:
 		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            : {_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}")
@@ -42,7 +42,7 @@ class Corner:
 
 		if IMAGE_DEBUG:
 			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)
+			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)
 
@@ -50,13 +50,13 @@ class Corner:
 		if CONSOLE_DEBUG:
 			print(f"{self.name}: {s}")
 
-	def process(self, image: np.array):
+	def process(self, _image: np.array):
-		image_anno = image.copy()
+		_image_anno = _image.copy()
-		crop_scaled = im.resize(image, IMG_SCALE_UP * image.shape[1], IMG_SCALE_UP * image.shape[0])
+		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
+		# Resize 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)
@@ -79,7 +79,7 @@ class Corner:
 
 		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)
+		self._debug(_image_anno, matcher_res, matcher_start_x, matcher_start_y, template_scaled, crop_scaled)
 
 		return matcher_bbox_local, matcher_res
 
@@ -114,7 +114,7 @@ class CornerTracker:
 		if CONSOLE_DEBUG:
 			print(f"{self.name}: {s}")
 
-	def _debug(self, image_anno, tracking_anno, matcher_bbox_local):
+	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)
@@ -124,20 +124,20 @@ class CornerTracker:
 		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)
+		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], self.color, 1)
+		cv2.rectangle(_image_anno, tracker_rect[0], tracker_rect[1], self.color, 1)
 
-	def init_reference_frame(self, image: np.array):
+	def init_reference_frame(self, _image: np.array):
 		print(f"Select tracking object")
-		bb = cv2.selectROI("Tracker Reference", image, False)
+		bb = cv2.selectROI("Tracker Reference", _image, False)
 		cv2.destroyWindow("Tracker Reference")
 		if bbox_center(bb) == (0, 0):
 			return False
 
 		self.tracking_ref_bb = bb
-		self.tracking_ref_img = image_crop(image, self.tracking_ref_bb)
+		self.tracking_ref_img = image_crop(_image, self.tracking_ref_bb)
 		self.tracking_ref_gray_img = cv2.cvtColor(self.tracking_ref_img, cv2.COLOR_BGR2GRAY)
 		self.corner_matcher_list = []
 
@@ -162,77 +162,80 @@ class CornerTracker:
 		cv2.destroyWindow("Matcher Reference")
 		print(f"Added {count} corners")
 
+		if len(corner_list) == 0:
+			return False
+
 		# Initialize tracker with first frame and bounding box
 		self.tracker = cv2.TrackerKCF().create()
-		self.tracker.init(image, self.tracking_ref_bb)
+		self.tracker.init(_image, self.tracking_ref_bb)
 
 		# Refine initial corners and store them as reference
-		self.corner_ref = self._corner_refine(cv2.cvtColor(image, cv2.COLOR_BGR2GRAY), corners=corner_list)
+		self.corner_ref = self._corner_refine(cv2.cvtColor(_image, cv2.COLOR_BGR2GRAY), corners=corner_list)
 
-		return len(corner_list) > 0
+		return True
 
-	def process(self, image: np.array, image_anno: 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")
 
 		# Update tracker
-		_ok, tracking_bb = self.tracker.update(image)
+		_ok, tracking_bb = self.tracker.update(_image)
 		if not _ok:
 			return None
 
-		tracking_img = image_crop(image.copy(), tracking_bb)
+		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, image_anno)
+		return self._match(_image, _image_anno)
 
-	def _match(self, image: np.array, image_anno: np.array):
+	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:
+		for _ct in self.corner_matcher_list:
-			matcher_bbox_local, matcher = ct.process(self.tracking_img)
+			matcher_bbox_local, matcher = _ct.process(self.tracking_img)
 			matcher_bbox = bbox_add_position(matcher_bbox_local, self.tracking_bb)
 
 			# Draw path
 			corner = bbox_center(matcher_bbox)
 			corners_raw.append(corner)
 
-			self._debug(image_anno, tracking_anno, matcher_bbox_local)
+			self._debug(_image_anno, tracking_anno, matcher_bbox_local)
 
 		# refine corners
-		corners_refined = self._corner_refine(cv2.cvtColor(image, cv2.COLOR_BGR2GRAY), corners_raw)
+		corners_refined = self._corner_refine(cv2.cvtColor(_image, cv2.COLOR_BGR2GRAY), corners_raw)
 
 		# store refined corners
-		i = 0
+		_i = 0
 		for corner in corners_refined:
-			ct = self.corner_matcher_list[i]
+			_ct = self.corner_matcher_list[_i]
-			ct.path_add(corner)
+			_ct.path_add(corner)
-			i += 1
+			_i += 1
 
 		# draw path
-		for ct in self.corner_matcher_list:
+		for _ct in self.corner_matcher_list:
-			for p in ct.path:
+			for p in _ct.path:
-				cv2.line(image_anno, bbox_round(p['from']), bbox_round(p['to']), (0, 0, 255), 1)
+				cv2.line(_image_anno, bbox_round(p['from']), bbox_round(p['to']), (0, 0, 255), 1)
 
 		distances = []
-		for i in range(0, corners_refined.shape[0]):
+		for _i in range(0, corners_refined.shape[0]):
-			distances.append((corners_refined[i] - self.corner_ref[i]))
+			distances.append((corners_refined[_i] - self.corner_ref[_i]))
 
-		mean_distance = np.mean(distances, axis=0)
+		_mean_distance = np.mean(distances, axis=0)
 
 		# show refined corners
-		for i in range(0, corners_refined.shape[0]):
+		for _i in range(0, corners_refined.shape[0]):
-			self._print(f" -- Corner Reference [{i}] {self.corner_ref[i]}")
+			self._print(f" -- Corner Reference [{i}] {self.corner_ref[_i]}")
-			self._print(f" -- Corner coarse    [{i}] {corners_raw[i]}")
+			self._print(f" -- Corner coarse    [{i}] {corners_raw[_i]}")
-			self._print(f" -- Corner fine      [{i}] {corners_refined[i]}")
+			self._print(f" -- Corner fine      [{i}] {corners_refined[_i]}")
-			self._print(f" -- Corner distance  [{i}] {corners_refined[i] - self.corner_ref[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.circle(_image_anno, (int(corners_refined[_i, 0]), int(corners_refined[_i, 1])), 4, (0, 255, 0))
 
 		if IMAGE_DEBUG:
 			cv2.imshow(f"{self.name}: Tracker: ", self.tracking_img)
 
-		return mean_distance
+		return _mean_distance
 
 	@staticmethod
 	def _corner_refine(src_gray, corners: np.array):
@@ -253,17 +256,14 @@ class CornerTracker:
 
 		# Write them down
 		corner_result = []
-		for i in range(corners_original.shape[0]):
+		for _i in range(corners_original.shape[0]):
-			corner = (corners_refined[i, 0, 0], corners_refined[i, 0, 1])
+			corner = (corners_refined[_i, 0, 0], corners_refined[_i, 0, 1])
 			corner_result.append(corner)
 
 		return np.array(corner_result)
 
 
 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
@@ -295,8 +295,8 @@ if __name__ == '__main__':
 				i = 0
 				for ct in tracker_list:
 					mean_distance = ct.process(image, image_anno)
-					cv2.putText(image_anno, f"Distance [{i}] : ({mean_distance[0]:+05.2f}, {mean_distance[1]:+05.2f})", (20, 25 + 25 * i),
+					cv2.putText(image_anno, f"Distance [{i}] : ({mean_distance[0]:+05.2f}, {mean_distance[1]:+05.2f})",
-								cv2.FONT_HERSHEY_SIMPLEX, 0.5, ct.color, 1)
+							(20, 25 + 25 * i), cv2.FONT_HERSHEY_SIMPLEX, 0.5, ct.color, 1)
 					i += 1
 				cv2.imshow(f"Image Anno", image_anno)
 			else: