diff --git a/ocv_corner_tracker.py b/ocv_corner_tracker.py index cc15e51..8c24af5 100644 --- a/ocv_corner_tracker.py +++ b/ocv_corner_tracker.py @@ -392,13 +392,12 @@ if __name__ == '__main__': ct_params = CornerTrackerParams(_scale=args.scale, _pre_track=args.track, _show_path=args.path) prj = project_load() + prj_loaded = False if prj is not None: ct_params = ct_params.from_dict(prj['params']) + prj_loaded = True # Let's go - colors = [(255, 0, 0), (0, 255, 0), (0, 0, 255), (255, 0, 255), (0, 255, 255), (255, 255, 255)] - tracker_count = 0 - tracker_list = [] if not video.isOpened(): print('Cannot open video file') sys.exit(1) @@ -408,23 +407,12 @@ if __name__ == '__main__': print('Cannot read video file') sys.exit(1) + colors = [(255, 0, 0), (0, 255, 0), (0, 0, 255), (255, 0, 255), (0, 255, 255), (255, 255, 255)] + tracker_count = 0 + tracker_list = [] 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 = [] + if prj_loaded: ct_params = CornerTrackerParams().from_dict(prj['params']) for tracker_settings in prj['trackers']: print(tracker_settings) @@ -433,6 +421,23 @@ if __name__ == '__main__': tracker_list.append(ct) tracker_count += 1 + else: + prj = {'params': ct_params.__dict__, 'trackers': []} + + print("Press 'a' for adding trackers, press 'd' to delete a tracker, press other key to continue") + k = cv2.waitKey(-1) & 0xff + while True: + print(f"Add 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 + tracker_count += 1 + project_save(prj) # rewind to frame 0