support project load/save
This commit is contained in:
+154
-58
@@ -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])
|
||||
rect_lower_right = (bb[0] + bb[2], bb[1] + bb[3])
|
||||
cv2.rectangle(_mask, rect_upper_left, rect_lower_right, 255, -1)
|
||||
|
||||
return True
|
||||
return bb
|
||||
|
||||
@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):
|
||||
# Draw mask
|
||||
def create_from_settings(self, _image: np.array, settings: CornerTrackerSettings):
|
||||
_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
|
||||
print(f"Select tracking object")
|
||||
bb = cv2.selectROI("Tracker Reference", _image, False)
|
||||
cv2.destroyWindow("Tracker Reference")
|
||||
if bbox_center(bb) == (0, 0):
|
||||
return False
|
||||
return _mask
|
||||
|
||||
tracking_img = image_crop(_image, bb)
|
||||
tracking_mask = image_crop(_mask, bb)
|
||||
def create_tracking(self, _image: np.array, _mask: np.array, tracking_bb: np.array):
|
||||
# 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,16 +408,32 @@ if __name__ == '__main__':
|
||||
print('Cannot read video file')
|
||||
sys.exit(1)
|
||||
|
||||
for tracker_count in range(0, 6):
|
||||
select_window = image.copy()
|
||||
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)
|
||||
if ok:
|
||||
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}")
|
||||
settings = ct.init_reference_frame(select_window)
|
||||
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)
|
||||
print(f"Number of active tracker: {len(tracker_list)}")
|
||||
else:
|
||||
break
|
||||
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]})
|
||||
|
||||
Reference in New Issue
Block a user