From f2f273f2d7f2d1776ebbbea1d00e6b12e2d6dc09 Mon Sep 17 00:00:00 2001 From: Jens Ahrensfeld Date: Fri, 19 Jul 2024 11:37:36 +0200 Subject: [PATCH] support project load/save --- .gitignore | 1 + ocv_corner_tracker.py | 212 ++++++++++++++++++++++++++++++------------ 2 files changed, 155 insertions(+), 58 deletions(-) diff --git a/.gitignore b/.gitignore index 644fb08..f72e6fd 100644 --- a/.gitignore +++ b/.gitignore @@ -1,2 +1,3 @@ __pycache__ *.csv +*.json diff --git a/ocv_corner_tracker.py b/ocv_corner_tracker.py index 6f11569..cc15e51 100644 --- a/ocv_corner_tracker.py +++ b/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]) - 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]})