support project load/save

.gitignore

@@ -1,2 +1,3 @@
 __pycache__
 *.csv
+*.json

ocv_corner_tracker.py

@@ -3,6 +3,7 @@ import numpy as np
 import argparse
 import sys
 import csv
+import json
 import matplotlib.pyplot as plt
 
 from util import image_crop, to_rect, bbox_add_position, bbox_center, bbox_round
@@ -87,13 +88,31 @@ class CornerTrackerParams:
 		self.pre_track = _pre_track
 		self.show_path = _show_path
 
+	def from_dict(self, params: dict):
+		self.__dict__.update(params)
+		return self
+
+
+class CornerTrackerSettings:
+	def __init__(self, _mask_bb=None, _tracking_bb=None):
+		self.mask_bb = _mask_bb
+		self.tracking_bb = _tracking_bb
+		self.matching_tpl_bb = []
+
+	def from_dict(self, settings: dict):
+		self.__dict__.update(settings)
+		return self
+
+	def append(self, _bb: np.array):
+		self.matching_tpl_bb.append(_bb)
+
 
 class CornerTracker:
 	def __init__(self, _params: CornerTrackerParams, color=(0, 255, 0), name: str = 'CornerTracker'):
 		self.params = _params
 		self.color = color
 		self.name = name
-		self.matching_tpl_bb = None
+		self.settings = None
 		self.tracking_bb = None
 		self.tracking_mask = None
 		self.tracker = None
@@ -127,11 +146,7 @@ class CornerTracker:
 		if bbox_center(bb) == (0, 0):
 			bb = (0, 0, _image.shape[1], _image.shape[0])
 
-		rect_upper_left = (bb[0], bb[1])
+		return bb
-		rect_lower_right = (bb[0] + bb[2], bb[1] + bb[3])
-		cv2.rectangle(_mask, rect_upper_left, rect_lower_right, 255, -1)
-
-		return True
 
 	@staticmethod
 	def mask_apply(_image: np.array, _mask: np.array):
@@ -147,55 +162,58 @@ class CornerTracker:
 		'''
 		return _result
 
-	def init_reference_frame(self, _image: np.array):
+	def create_from_settings(self, _image: np.array, settings: CornerTrackerSettings):
-		# Draw mask
+		_mask = CornerTracker.create_mask(_image, settings.mask_bb)
+		image_processed_local, _ = self.create_tracking(_image, _mask, settings.tracking_bb)
+		self.create_matching(image_processed_local, settings.tracking_bb, settings.matching_tpl_bb)
+
+		return self.create_matching(image_processed_local, settings.tracking_bb, settings.matching_tpl_bb)
+
+	@staticmethod
+	def create_mask(_image: np.array, mask_bb: np.array):
+		# Init mask
 		_mask = CornerTracker.mask_init(_image)
-		CornerTracker.mask_draw(_image, _mask)
+		rect_upper_left = (mask_bb[0], mask_bb[1])
+		rect_lower_right = (mask_bb[0] + mask_bb[2], mask_bb[1] + mask_bb[3])
+		# Mask cut out
+		cv2.rectangle(_mask, rect_upper_left, rect_lower_right, 255, -1)
 
-		# Select search area
+		return _mask
-		print(f"Select tracking object")
-		bb = cv2.selectROI("Tracker Reference", _image, False)
-		cv2.destroyWindow("Tracker Reference")
-		if bbox_center(bb) == (0, 0):
-			return False
 
-		tracking_img = image_crop(_image, bb)
+	def create_tracking(self, _image: np.array, _mask: np.array, tracking_bb: np.array):
-		tracking_mask = image_crop(_mask, bb)
+		# Init search area
+		tracking_img = image_crop(_image, tracking_bb)
+		tracking_mask = image_crop(_mask, tracking_bb)
 		CornerTracker.mask_apply(tracking_img, tracking_mask)
 		image_processed_local = CornerTracker.image_process(tracking_img)
 
-		self.corner_matcher_list = []
-		self.tracking_bb = bb
-		self.tracking_mask = tracking_mask
-
-		corner_list = []
-		count = 1
-		while True:
-			print(f"Draw Corner #{count}")
-			matching_tpl_bb = cv2.selectROI("Matcher Reference", tracking_img, fromCenter=True, showCrosshair=True)
-			if bbox_center(matching_tpl_bb) == (0, 0):
-				break
-
-			matching_tpl_img = image_crop(image_processed_local, 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")
-
-			count += 1
-
-		cv2.destroyWindow("Matcher Reference")
-		print(f"Added {count} corners")
-
-		if len(corner_list) == 0:
-			return False
-
 		# Initialize tracker with first frame and bounding box
 		if self.params.pre_track:
 			self.tracker = cv2.TrackerKCF.create()
-			self.tracker.init(_image, self.tracking_bb)
+			self.tracker.init(_image, tracking_bb)
+
+		self.tracking_bb = tracking_bb
+		self.tracking_mask = image_crop(_mask, tracking_bb)
+
+		return image_processed_local, tracking_img
+
+	def create_matching(self, image_processed_local: np.array, tracking_bb: np.array, matching_tpl_bb_list: list):
+		# Create corner matcher
+		self.corner_matcher_list = []
+		corner_list = []
+		count = 1
+		for matching_tpl_bb in matching_tpl_bb_list:
+			if bbox_center(matching_tpl_bb) == (0, 0):
+				break
+			corner_list.append(bbox_center(matching_tpl_bb))
+			matching_tpl_img = image_crop(image_processed_local, matching_tpl_bb)
+			self.corner_matcher_list.append(Corner(matching_tpl_img, matching_tpl_bb, name=f"Corner-{count}"))
+
+			print(f"Corner {count} added")
+			count += 1
+
+		if len(corner_list) == 0:
+			return False
 
 		# Refine initial corners and store them as reference
 		corners_local = self._corner_refine(image_processed_local, corners=corner_list)
@@ -203,11 +221,53 @@ class CornerTracker:
 		# 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])
+			corner = (corner_local[0] + tracking_bb[0], corner_local[1] + tracking_bb[1])
 			self.corner_ref.append(corner)
 
 		return True
 
+	def init_reference_frame(self, _image: np.array):
+		settings = CornerTrackerSettings()
+		# Draw mask
+		_mask = CornerTracker.mask_init(_image)
+
+		# Init mask
+		print(f"Draw mask object")
+		_mask_bb = cv2.selectROI("Tracker Reference", _image, False)
+		cv2.destroyWindow("Tracker Reference")
+		if bbox_center(_mask_bb) == (0, 0):
+			_mask_bb = (0, 0, _image.shape[1], _image.shape[0])
+
+		settings.mask_bb = _mask_bb
+		_mask = self.create_mask(_image, _mask_bb)
+
+		# Select search area
+		print(f"Select tracking object")
+		tracking_bb = cv2.selectROI("Tracker Reference", _image, False)
+		cv2.destroyWindow("Tracker Reference")
+		if bbox_center(tracking_bb) == (0, 0):
+			return None
+
+		settings.tracking_bb = tracking_bb
+		image_processed_local, tracking_img = self.create_tracking(_image, _mask, tracking_bb)
+
+		count = 1
+		while True:
+			print(f"Draw Corner #{count}")
+			matching_tpl_bb = cv2.selectROI("Matcher Reference", tracking_img, fromCenter=True, showCrosshair=True)
+			if bbox_center(matching_tpl_bb) == (0, 0):
+				break
+
+			settings.matching_tpl_bb.append(matching_tpl_bb)
+			print(f"Corner {count} added")
+			count += 1
+
+		cv2.destroyWindow("Matcher Reference")
+		print(f"Added {count} corners")
+
+		if self.create_matching(image_processed_local, settings.tracking_bb, settings.matching_tpl_bb):
+			return settings
+
 	def process(self, _image: np.array, _image_anno: np.array):
 		# Update tracker bb
 		if self.tracker is not None:
@@ -298,6 +358,22 @@ class CornerTracker:
 		return np.array(corner_result)
 
 
+def project_load(path: str = "./"):
+	_settings = None
+	try:
+		with open(path + "ocv_tracker_prj.json", "r") as fp:
+			_settings = json.load(fp)
+	except Exception:
+		pass
+
+	return _settings
+
+
+def project_save(settings: dict, path: str = "./"):
+	with open(path + "ocv_tracker_prj.json", "w") as fp:
+		json.dump(settings, fp, indent=4)
+
+
 if __name__ == '__main__':
 	# Parse command line args
 	parser = argparse.ArgumentParser(description="Gearbox Tracker")
@@ -313,7 +389,11 @@ if __name__ == '__main__':
 	video = cv2.VideoCapture(args.filename)
 
 	# Fill CornerTrackerParams from  args
-	ct_params = CornerTrackerParams(_scale=np.float32(args.scale), _pre_track=args.track, _show_path=args.path)
+	ct_params = CornerTrackerParams(_scale=args.scale, _pre_track=args.track, _show_path=args.path)
+
+	prj = project_load()
+	if prj is not None:
+		ct_params = ct_params.from_dict(prj['params'])
 
 	# Let's go
 	colors = [(255, 0, 0), (0, 255, 0), (0, 0, 255), (255, 0, 255), (0, 255, 255), (255, 255, 255)]
@@ -328,17 +408,33 @@ if __name__ == '__main__':
 		print('Cannot read video file')
 		sys.exit(1)
 
-	for tracker_count in range(0, 6):
 	select_window = image.copy()
+	if prj is None:
+		prj = {'params': ct_params.__dict__, 'trackers': []}
+		for tracker_count in range(0, 6):
 			print(f"Draw tracker #{tracker_count}")
 			ct = CornerTracker(ct_params, colors[tracker_count], name=f"Tracker-{tracker_count}")
-		ok = ct.init_reference_frame(select_window)
+			settings = ct.init_reference_frame(select_window)
-		if ok:
+			if settings is not None:
+				prj['trackers'].append(settings.__dict__)
 				tracker_list.append(ct)
 				print(f"Number of active tracker: {len(tracker_list)}")
 			else:
 				break
 
+	else:
+		tracker_count = 0
+		tracker_list = []
+		ct_params = CornerTrackerParams().from_dict(prj['params'])
+		for tracker_settings in prj['trackers']:
+			print(tracker_settings)
+			ct = CornerTracker(ct_params, colors[tracker_count], name=f"Tracker-{tracker_count}")
+			ct.create_from_settings(select_window, CornerTrackerSettings().from_dict(tracker_settings))
+			tracker_list.append(ct)
+			tracker_count += 1
+
+	project_save(prj)
+
 	# rewind to frame 0
 	video.set(cv2.CAP_PROP_POS_FRAMES, 0)
 
@@ -365,7 +461,7 @@ if __name__ == '__main__':
 		for ct in tracker_list:
 			mean_distance = ct.process(image, image_anno)
 			if mean_distance is not None:
-				scale = ct_params.scale
+				scale = np.float32(ct_params.scale)
 				scaled_distance = (scale*mean_distance[0], scale*mean_distance[1])
 				if args.csv:
 					result_csv[i].append({f'dx': scaled_distance[0], 'dy': scaled_distance[1]})