def __init__(self, params): self.params = params self.pause_mode = False self.step = False self.visdom = None if self.params.debug > 0 and self.params.visdom_info.get( 'use_visdom', True): try: self.visdom = Visdom(self.params.debug, { 'handler': self.visdom_ui_handler, 'win_id': 'Tracking' }, visdom_info=self.params.visdom_info) # Show help help_text = 'You can pause/unpause the tracker by pressing ''space'' with the ''Tracking'' window ' \ 'selected. During paused mode, you can track for one frame by pressing the right arrow key.' \ 'To enable/disable plotting of a data block, tick/untick the corresponding entry in ' \ 'block list.' self.visdom.register(help_text, 'text', 1, 'Help') except: time.sleep(0.5) print( '!!! WARNING: Visdom could not start, so using matplotlib visualization instead !!!\n' '!!! Start Visdom in a separate terminal window by typing \'visdom\' !!!' )
class Tracker: """Wraps the tracker for evaluation and running purposes. args: name: Name of tracking method. parameter_name: Name of parameter file. run_id: The run id. display_name: Name to be displayed in the result plots. """ def __init__(self, name: str, parameter_name: str, run_id: int = None, display_name: str = None): assert run_id is None or isinstance(run_id, int) self.name = name self.parameter_name = parameter_name self.run_id = run_id self.display_name = display_name env = env_settings() if self.run_id is None: self.results_dir = '{}/{}/{}'.format(env.results_path, self.name, self.parameter_name) self.segmentation_dir = '{}/{}/{}'.format(env.segmentation_path, self.name, self.parameter_name) else: self.results_dir = '{}/{}/{}_{:03d}'.format(env.results_path, self.name, self.parameter_name, self.run_id) self.segmentation_dir = '{}/{}/{}_{:03d}'.format(env.segmentation_path, self.name, self.parameter_name, self.run_id) tracker_module_abspath = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', 'tracker', self.name)) if os.path.isdir(tracker_module_abspath): tracker_module = importlib.import_module('pytracking.tracker.{}'.format(self.name)) self.tracker_class = tracker_module.get_tracker_class() else: self.tracker_class = None self.visdom = None def _init_visdom(self, visdom_info, debug): visdom_info = {} if visdom_info is None else visdom_info self.pause_mode = False self.step = False if debug > 0 and visdom_info.get('use_visdom', True): try: self.visdom = Visdom(debug, {'handler': self._visdom_ui_handler, 'win_id': 'Tracking'}, visdom_info=visdom_info) # Show help help_text = 'You can pause/unpause the tracker by pressing ''space'' with the ''Tracking'' window ' \ 'selected. During paused mode, you can track for one frame by pressing the right arrow key.' \ 'To enable/disable plotting of a data block, tick/untick the corresponding entry in ' \ 'block list.' self.visdom.register(help_text, 'text', 1, 'Help') except: time.sleep(0.5) print('!!! WARNING: Visdom could not start, so using matplotlib visualization instead !!!\n' '!!! Start Visdom in a separate terminal window by typing \'visdom\' !!!') def _visdom_ui_handler(self, data): if data['event_type'] == 'KeyPress': if data['key'] == ' ': self.pause_mode = not self.pause_mode elif data['key'] == 'ArrowRight' and self.pause_mode: self.step = True def create_tracker(self, params): tracker = self.tracker_class(params) tracker.visdom = self.visdom return tracker def run_sequence(self, seq, visualization=None, debug=None, visdom_info=None, multiobj_mode=None): """Run tracker on sequence. args: seq: Sequence to run the tracker on. visualization: Set visualization flag (None means default value specified in the parameters). debug: Set debug level (None means default value specified in the parameters). visdom_info: Visdom info. multiobj_mode: Which mode to use for multiple objects. """ params = self.get_parameters() visualization_ = visualization debug_ = debug if debug is None: debug_ = getattr(params, 'debug', 0) if visualization is None: if debug is None: visualization_ = getattr(params, 'visualization', False) else: visualization_ = True if debug else False params.visualization = visualization_ params.debug = debug_ self._init_visdom(visdom_info, debug_) if visualization_ and self.visdom is None: self.init_visualization() # Get init information init_info = seq.init_info() is_single_object = not seq.multiobj_mode if multiobj_mode is None: multiobj_mode = getattr(params, 'multiobj_mode', getattr(self.tracker_class, 'multiobj_mode', 'default')) if multiobj_mode == 'default' or is_single_object: tracker = self.create_tracker(params) elif multiobj_mode == 'parallel': tracker = MultiObjectWrapper(self.tracker_class, params, self.visdom) else: raise ValueError('Unknown multi object mode {}'.format(multiobj_mode)) output = self._track_sequence(tracker, seq, init_info) return output def _track_sequence(self, tracker, seq, init_info): # Define outputs # Each field in output is a list containing tracker prediction for each frame. # In case of single object tracking mode: # target_bbox[i] is the predicted bounding box for frame i # time[i] is the processing time for frame i # segmentation[i] is the segmentation mask for frame i (numpy array) # In case of multi object tracking mode: # target_bbox[i] is an OrderedDict, where target_bbox[i][obj_id] is the predicted box for target obj_id in # frame i # time[i] is either the processing time for frame i, or an OrderedDict containing processing times for each # object in frame i # segmentation[i] is the multi-label segmentation mask for frame i (numpy array) output = {'target_bbox': [], 'time': [], 'segmentation': []} def _store_outputs(tracker_out: dict, defaults=None): defaults = {} if defaults is None else defaults for key in output.keys(): val = tracker_out.get(key, defaults.get(key, None)) if key in tracker_out or val is not None: output[key].append(val) # Initialize image = self._read_image(seq.frames[0]) if tracker.params.visualization and self.visdom is None: self.visualize(image, init_info.get('init_bbox')) start_time = time.time() out = tracker.initialize(image, init_info) if out is None: out = {} prev_output = OrderedDict(out) init_default = {'target_bbox': init_info.get('init_bbox'), 'time': time.time() - start_time, 'segmentation': init_info.get('init_mask')} _store_outputs(out, init_default) for frame_num, frame_path in enumerate(seq.frames[1:], start=1): while True: if not self.pause_mode: break elif self.step: self.step = False break else: time.sleep(0.1) image = self._read_image(frame_path) start_time = time.time() info = seq.frame_info(frame_num) info['previous_output'] = prev_output out = tracker.track(image, info) prev_output = OrderedDict(out) _store_outputs(out, {'time': time.time() - start_time}) segmentation = out['segmentation'] if 'segmentation' in out else None if self.visdom is not None: tracker.visdom_draw_tracking(image, out['target_bbox'], segmentation) elif tracker.params.visualization: self.visualize(image, out['target_bbox'], segmentation) for key in ['target_bbox', 'segmentation']: if key in output and len(output[key]) <= 1: output.pop(key) return output def run_video(self, videofilepath, optional_box=None, debug=None, visdom_info=None, save_results=False): """Run the tracker with the vieofile. args: debug: Debug level. """ params = self.get_parameters() debug_ = debug if debug is None: debug_ = getattr(params, 'debug', 0) params.debug = debug_ params.tracker_name = self.name params.param_name = self.parameter_name self._init_visdom(visdom_info, debug_) multiobj_mode = getattr(params, 'multiobj_mode', getattr(self.tracker_class, 'multiobj_mode', 'default')) if multiobj_mode == 'default': tracker = self.create_tracker(params) if hasattr(tracker, 'initialize_features'): tracker.initialize_features() elif multiobj_mode == 'parallel': tracker = MultiObjectWrapper(self.tracker_class, params, self.visdom, fast_load=True) else: raise ValueError('Unknown multi object mode {}'.format(multiobj_mode)) assert os.path.isfile(videofilepath), "Invalid param {}".format(videofilepath) ", videofilepath must be a valid videofile" output_boxes = [] cap = cv.VideoCapture(videofilepath) display_name = 'Display: ' + tracker.params.tracker_name cv.namedWindow(display_name, cv.WINDOW_NORMAL | cv.WINDOW_KEEPRATIO) cv.resizeWindow(display_name, 960, 720) success, frame = cap.read() cv.imshow(display_name, frame) def _build_init_info(box): return {'init_bbox': OrderedDict({1: box}), 'init_object_ids': [1, ], 'object_ids': [1, ], 'sequence_object_ids': [1, ]} if success is not True: print("Read frame from {} failed.".format(videofilepath)) exit(-1) if optional_box is not None: assert isinstance(optional_box, (list, tuple)) assert len(optional_box) == 4, "valid box's foramt is [x,y,w,h]" tracker.initialize(frame, _build_init_info(optional_box)) output_boxes.append(optional_box) else: while True: # cv.waitKey() frame_disp = frame.copy() cv.putText(frame_disp, 'Select target ROI and press ENTER', (20, 30), cv.FONT_HERSHEY_COMPLEX_SMALL, 1.5, (0, 0, 0), 1) x, y, w, h = cv.selectROI(display_name, frame_disp, fromCenter=False) init_state = [x, y, w, h] tracker.initialize(frame, _build_init_info(init_state)) output_boxes.append(init_state) break while True: ret, frame = cap.read() if frame is None: break frame_disp = frame.copy() # Draw box out = tracker.track(frame) state = [int(s) for s in out['target_bbox'][1]] output_boxes.append(state) cv.rectangle(frame_disp, (state[0], state[1]), (state[2] + state[0], state[3] + state[1]), (0, 255, 0), 5) font_color = (0, 0, 0) cv.putText(frame_disp, 'Tracking!', (20, 30), cv.FONT_HERSHEY_COMPLEX_SMALL, 1, font_color, 1) cv.putText(frame_disp, 'Press r to reset', (20, 55), cv.FONT_HERSHEY_COMPLEX_SMALL, 1, font_color, 1) cv.putText(frame_disp, 'Press q to quit', (20, 80), cv.FONT_HERSHEY_COMPLEX_SMALL, 1, font_color, 1) # Display the resulting frame cv.imshow(display_name, frame_disp) key = cv.waitKey(1) if key == ord('q'): break elif key == ord('r'): ret, frame = cap.read() frame_disp = frame.copy() cv.putText(frame_disp, 'Select target ROI and press ENTER', (20, 30), cv.FONT_HERSHEY_COMPLEX_SMALL, 1.5, (0, 0, 0), 1) cv.imshow(display_name, frame_disp) x, y, w, h = cv.selectROI(display_name, frame_disp, fromCenter=False) init_state = [x, y, w, h] tracker.initialize(frame, _build_init_info(init_state)) output_boxes.append(init_state) # When everything done, release the capture cap.release() cv.destroyAllWindows() if save_results: if not os.path.exists(self.results_dir): os.makedirs(self.results_dir) video_name = Path(videofilepath).stem base_results_path = os.path.join(self.results_dir, 'video_{}'.format(video_name)) tracked_bb = np.array(output_boxes).astype(int) bbox_file = '{}.txt'.format(base_results_path) np.savetxt(bbox_file, tracked_bb, delimiter='\t', fmt='%d') def run_webcam(self, debug=None, visdom_info=None): """Run the tracker with the webcam. args: debug: Debug level. """ params = self.get_parameters() debug_ = debug if debug is None: debug_ = getattr(params, 'debug', 0) params.debug = debug_ params.tracker_name = self.name params.param_name = self.parameter_name self._init_visdom(visdom_info, debug_) multiobj_mode = getattr(params, 'multiobj_mode', getattr(self.tracker_class, 'multiobj_mode', 'default')) if multiobj_mode == 'default': tracker = self.create_tracker(params) elif multiobj_mode == 'parallel': tracker = MultiObjectWrapper(self.tracker_class, params, self.visdom, fast_load=True) else: raise ValueError('Unknown multi object mode {}'.format(multiobj_mode)) class UIControl: def __init__(self): self.mode = 'init' # init, select, track self.target_tl = (-1, -1) self.target_br = (-1, -1) self.new_init = False def mouse_callback(self, event, x, y, flags, param): if event == cv.EVENT_LBUTTONDOWN and self.mode == 'init': self.target_tl = (x, y) self.target_br = (x, y) self.mode = 'select' elif event == cv.EVENT_MOUSEMOVE and self.mode == 'select': self.target_br = (x, y) elif event == cv.EVENT_LBUTTONDOWN and self.mode == 'select': self.target_br = (x, y) self.mode = 'init' self.new_init = True def get_tl(self): return self.target_tl if self.target_tl[0] < self.target_br[0] else self.target_br def get_br(self): return self.target_br if self.target_tl[0] < self.target_br[0] else self.target_tl def get_bb(self): tl = self.get_tl() br = self.get_br() bb = [min(tl[0], br[0]), min(tl[1], br[1]), abs(br[0] - tl[0]), abs(br[1] - tl[1])] return bb ui_control = UIControl() cap = cv.VideoCapture(0) display_name = 'Display: ' + self.name cv.namedWindow(display_name, cv.WINDOW_NORMAL | cv.WINDOW_KEEPRATIO) cv.resizeWindow(display_name, 960, 720) cv.setMouseCallback(display_name, ui_control.mouse_callback) next_object_id = 1 sequence_object_ids = [] prev_output = OrderedDict() while True: # Capture frame-by-frame ret, frame = cap.read() frame_disp = frame.copy() info = OrderedDict() info['previous_output'] = prev_output if ui_control.new_init: ui_control.new_init = False init_state = ui_control.get_bb() info['init_object_ids'] = [next_object_id, ] info['init_bbox'] = OrderedDict({next_object_id: init_state}) sequence_object_ids.append(next_object_id) next_object_id += 1 # Draw box if ui_control.mode == 'select': cv.rectangle(frame_disp, ui_control.get_tl(), ui_control.get_br(), (255, 0, 0), 2) if len(sequence_object_ids) > 0: info['sequence_object_ids'] = sequence_object_ids out = tracker.track(frame, info) prev_output = OrderedDict(out) if 'segmentation' in out: frame_disp = overlay_mask(frame_disp, out['segmentation']) if 'target_bbox' in out: for obj_id, state in out['target_bbox'].items(): state = [int(s) for s in state] cv.rectangle(frame_disp, (state[0], state[1]), (state[2] + state[0], state[3] + state[1]), _tracker_disp_colors[obj_id], 5) # Put text font_color = (0, 0, 0) cv.putText(frame_disp, 'Select target', (20, 30), cv.FONT_HERSHEY_COMPLEX_SMALL, 1, font_color, 1) cv.putText(frame_disp, 'Press r to reset', (20, 55), cv.FONT_HERSHEY_COMPLEX_SMALL, 1, font_color, 1) cv.putText(frame_disp, 'Press q to quit', (20, 85), cv.FONT_HERSHEY_COMPLEX_SMALL, 1, font_color, 1) # Display the resulting frame cv.imshow(display_name, frame_disp) key = cv.waitKey(1) if key == ord('q'): break elif key == ord('r'): next_object_id = 1 sequence_object_ids = [] prev_output = OrderedDict() info = OrderedDict() info['object_ids'] = [] info['init_object_ids'] = [] info['init_bbox'] = OrderedDict() tracker.initialize(frame, info) ui_control.mode = 'init' # When everything done, release the capture cap.release() cv.destroyAllWindows() def run_vot2020(self, debug=None, visdom_info=None): params = self.get_parameters() params.tracker_name = self.name params.param_name = self.parameter_name params.run_id = self.run_id debug_ = debug if debug is None: debug_ = getattr(params, 'debug', 0) if debug is None: visualization_ = getattr(params, 'visualization', False) else: visualization_ = True if debug else False params.visualization = visualization_ params.debug = debug_ self._init_visdom(visdom_info, debug_) tracker = self.create_tracker(params) tracker.initialize_features() output_segmentation = tracker.predicts_segmentation_mask() import pytracking.evaluation.vot2020 as vot def _convert_anno_to_list(vot_anno): vot_anno = [vot_anno[0], vot_anno[1], vot_anno[2], vot_anno[3]] return vot_anno def _convert_image_path(image_path): return image_path """Run tracker on VOT.""" if output_segmentation: handle = vot.VOT("mask") else: handle = vot.VOT("rectangle") vot_anno = handle.region() image_path = handle.frame() if not image_path: return image_path = _convert_image_path(image_path) image = self._read_image(image_path) if output_segmentation: vot_anno_mask = vot.make_full_size(vot_anno, (image.shape[1], image.shape[0])) bbox = masks_to_bboxes(torch.from_numpy(vot_anno_mask), fmt='t').squeeze().tolist() else: bbox = _convert_anno_to_list(vot_anno) vot_anno_mask = None out = tracker.initialize(image, {'init_mask': vot_anno_mask, 'init_bbox': bbox}) if out is None: out = {} prev_output = OrderedDict(out) # Track while True: image_path = handle.frame() if not image_path: break image_path = _convert_image_path(image_path) image = self._read_image(image_path) info = OrderedDict() info['previous_output'] = prev_output out = tracker.track(image, info) prev_output = OrderedDict(out) if output_segmentation: pred = out['segmentation'].astype(np.uint8) else: state = out['target_bbox'] pred = vot.Rectangle(*state) handle.report(pred, 1.0) segmentation = out['segmentation'] if 'segmentation' in out else None if self.visdom is not None: tracker.visdom_draw_tracking(image, out['target_bbox'], segmentation) elif tracker.params.visualization: self.visualize(image, out['target_bbox'], segmentation) def run_vot(self, debug=None, visdom_info=None): params = self.get_parameters() params.tracker_name = self.name params.param_name = self.parameter_name params.run_id = self.run_id debug_ = debug if debug is None: debug_ = getattr(params, 'debug', 0) if debug is None: visualization_ = getattr(params, 'visualization', False) else: visualization_ = True if debug else False params.visualization = visualization_ params.debug = debug_ self._init_visdom(visdom_info, debug_) tracker = self.create_tracker(params) tracker.initialize_features() import pytracking.evaluation.vot as vot def _convert_anno_to_list(vot_anno): vot_anno = [vot_anno[0][0][0], vot_anno[0][0][1], vot_anno[0][1][0], vot_anno[0][1][1], vot_anno[0][2][0], vot_anno[0][2][1], vot_anno[0][3][0], vot_anno[0][3][1]] return vot_anno def _convert_image_path(image_path): image_path_new = image_path[20:- 2] return "".join(image_path_new) """Run tracker on VOT.""" handle = vot.VOT("polygon") vot_anno_polygon = handle.region() vot_anno_polygon = _convert_anno_to_list(vot_anno_polygon) init_state = convert_vot_anno_to_rect(vot_anno_polygon, tracker.params.vot_anno_conversion_type) image_path = handle.frame() if not image_path: return image_path = _convert_image_path(image_path) image = self._read_image(image_path) tracker.initialize(image, {'init_bbox': init_state}) # Track while True: image_path = handle.frame() if not image_path: break image_path = _convert_image_path(image_path) image = self._read_image(image_path) out = tracker.track(image) state = out['target_bbox'] handle.report(vot.Rectangle(state[0], state[1], state[2], state[3])) segmentation = out['segmentation'] if 'segmentation' in out else None if self.visdom is not None: tracker.visdom_draw_tracking(image, out['target_bbox'], segmentation) elif tracker.params.visualization: self.visualize(image, out['target_bbox'], segmentation) def get_parameters(self): """Get parameters.""" param_module = importlib.import_module('pytracking.parameter.{}.{}'.format(self.name, self.parameter_name)) params = param_module.parameters() return params def init_visualization(self): self.pause_mode = False self.fig, self.ax = plt.subplots(1) self.fig.canvas.mpl_connect('key_press_event', self.press) plt.tight_layout() def visualize(self, image, state, segmentation=None): self.ax.cla() self.ax.imshow(image) if segmentation is not None: self.ax.imshow(segmentation, alpha=0.5) if isinstance(state, (OrderedDict, dict)): boxes = [v for k, v in state.items()] else: boxes = (state,) for i, box in enumerate(boxes, start=1): col = _tracker_disp_colors[i] col = [float(c) / 255.0 for c in col] rect = patches.Rectangle((box[0], box[1]), box[2], box[3], linewidth=1, edgecolor=col, facecolor='none') self.ax.add_patch(rect) if getattr(self, 'gt_state', None) is not None: gt_state = self.gt_state rect = patches.Rectangle((gt_state[0], gt_state[1]), gt_state[2], gt_state[3], linewidth=1, edgecolor='g', facecolor='none') self.ax.add_patch(rect) self.ax.set_axis_off() self.ax.axis('equal') draw_figure(self.fig) if self.pause_mode: keypress = False while not keypress: keypress = plt.waitforbuttonpress() def reset_tracker(self): pass def press(self, event): if event.key == 'p': self.pause_mode = not self.pause_mode print("Switching pause mode!") elif event.key == 'r': self.reset_tracker() print("Resetting target pos to gt!") def _read_image(self, image_file: str): im = cv.imread(image_file) return cv.cvtColor(im, cv.COLOR_BGR2RGB)
class BaseTracker: """Base class for all trackers.""" def visdom_ui_handler(self, data): if data['event_type'] == 'KeyPress': if data['key'] == ' ': self.pause_mode = not self.pause_mode elif data['key'] == 'ArrowRight' and self.pause_mode: self.step = True def __init__(self, params): self.params = params self.pause_mode = False self.step = False self.visdom = None if self.params.debug > 0 and self.params.visdom_info.get( 'use_visdom', True): try: self.visdom = Visdom(self.params.debug, { 'handler': self.visdom_ui_handler, 'win_id': 'Tracking' }, visdom_info=self.params.visdom_info) # Show help help_text = 'You can pause/unpause the tracker by pressing ''space'' with the ''Tracking'' window ' \ 'selected. During paused mode, you can track for one frame by pressing the right arrow key.' \ 'To enable/disable plotting of a data block, tick/untick the corresponding entry in ' \ 'block list.' self.visdom.register(help_text, 'text', 1, 'Help') except: time.sleep(0.5) print( '!!! WARNING: Visdom could not start, so using matplotlib visualization instead !!!\n' '!!! Start Visdom in a separate terminal window by typing \'visdom\' !!!' ) def initialize(self, image, info: dict) -> dict: """Overload this function in your tracker. This should initialize the model.""" raise NotImplementedError def track(self, image) -> dict: """Overload this function in your tracker. This should track in the frame and update the model.""" raise NotImplementedError def track_sequence(self, sequence): """Run tracker on a sequence.""" output = {'target_bbox': [], 'time': []} def _store_outputs(tracker_out: dict, defaults=None): defaults = {} if defaults is None else defaults for key in tracker_out.keys(): if key not in output: raise RuntimeError('Unknown output from tracker.') for key in output.keys(): val = tracker_out.get(key, defaults.get(key, None)) if val is not None: output[key].append(val) # Initialize image = self._read_image(sequence.frames[0]) if self.params.visualization and self.visdom is None: self.init_visualization() self.visualize(image, sequence.get('init_bbox'), sequence.name, 0) start_time = time.time() out = self.initialize(image, sequence.init_info()) if out is None: out = {} _store_outputs( out, { 'target_bbox': sequence.get('init_bbox'), 'time': time.time() - start_time }) if self.visdom is not None: self.visdom.register((image, sequence.get('init_bbox')), 'Tracking', 1, 'Tracking') # Track for i, frame in enumerate(sequence.frames[1:]): while True: if not self.pause_mode: break elif self.step: self.step = False break else: time.sleep(0.1) image = self._read_image(frame) start_time = time.time() out = self.track(image) _store_outputs(out, {'time': time.time() - start_time}) if self.visdom is not None: self.visdom.register((image, out['target_bbox']), 'Tracking', 1, 'Tracking') elif self.params.visualization: self.visualize(image, out['target_bbox'], sequence.name, i + 1) return output def track_videofile(self, videofilepath, optional_box=None): """Run track with a video file input.""" assert os.path.isfile(videofilepath), "Invalid param {}".format( videofilepath) ", videofilepath must be a valid videofile" if hasattr(self, 'initialize_features'): self.initialize_features() cap = cv.VideoCapture(videofilepath) display_name = 'Display: ' + self.params.tracker_name cv.namedWindow(display_name, cv.WINDOW_NORMAL | cv.WINDOW_KEEPRATIO) cv.resizeWindow(display_name, 960, 720) success, frame = cap.read() cv.imshow(display_name, frame) if success is not True: print("Read frame from {} failed.".format(videofilepath)) exit(-1) if optional_box is not None: assert isinstance(optional_box, list, tuple) assert len(optional_box) == 4, "valid box's foramt is [x,y,w,h]" self.initialize(frame, {'init_bbox': optional_box}) else: while True: # cv.waitKey() frame_disp = frame.copy() cv.putText(frame_disp, 'Select target ROI and press ENTER', (20, 30), cv.FONT_HERSHEY_COMPLEX_SMALL, 1.5, (0, 0, 0), 1) x, y, w, h = cv.selectROI(display_name, frame_disp, fromCenter=False) init_state = [x, y, w, h] self.initialize(frame, {'init_bbox': init_state}) break while True: ret, frame = cap.read() if frame is None: return frame_disp = frame.copy() # Draw box out = self.track(frame) state = [int(s) for s in out['target_bbox']] cv.rectangle(frame_disp, (state[0], state[1]), (state[2] + state[0], state[3] + state[1]), (0, 255, 0), 5) font_color = (0, 0, 0) cv.putText(frame_disp, 'Tracking!', (20, 30), cv.FONT_HERSHEY_COMPLEX_SMALL, 1, font_color, 1) cv.putText(frame_disp, 'Press r to reset', (20, 55), cv.FONT_HERSHEY_COMPLEX_SMALL, 1, font_color, 1) cv.putText(frame_disp, 'Press q to quit', (20, 80), cv.FONT_HERSHEY_COMPLEX_SMALL, 1, font_color, 1) # Display the resulting frame cv.imshow(display_name, frame_disp) key = cv.waitKey(1) if key == ord('q'): break elif key == ord('r'): ret, frame = cap.read() frame_disp = frame.copy() cv.putText(frame_disp, 'Select target ROI and press ENTER', (20, 30), cv.FONT_HERSHEY_COMPLEX_SMALL, 1.5, (0, 0, 0), 1) cv.imshow(display_name, frame_disp) x, y, w, h = cv.selectROI(display_name, frame_disp, fromCenter=False) init_state = [x, y, w, h] self.initialize(frame, {'init_bbox': init_state}) # When everything done, release the capture cap.release() cv.destroyAllWindows() def track_webcam(self): """Run tracker with webcam.""" class UIControl: def __init__(self): self.mode = 'init' # init, select, track self.target_tl = (-1, -1) self.target_br = (-1, -1) self.mode_switch = False def mouse_callback(self, event, x, y, flags, param): if event == cv.EVENT_LBUTTONDOWN and self.mode == 'init': self.target_tl = (x, y) self.target_br = (x, y) self.mode = 'select' self.mode_switch = True elif event == cv.EVENT_MOUSEMOVE and self.mode == 'select': self.target_br = (x, y) elif event == cv.EVENT_LBUTTONDOWN and self.mode == 'select': self.target_br = (x, y) self.mode = 'track' self.mode_switch = True def get_tl(self): return self.target_tl if self.target_tl[0] < self.target_br[ 0] else self.target_br def get_br(self): return self.target_br if self.target_tl[0] < self.target_br[ 0] else self.target_tl def get_bb(self): tl = self.get_tl() br = self.get_br() bb = [ min(tl[0], br[0]), min(tl[1], br[1]), abs(br[0] - tl[0]), abs(br[1] - tl[1]) ] return bb ui_control = UIControl() cap = cv.VideoCapture(0) display_name = 'Display: ' + self.params.tracker_name cv.namedWindow(display_name, cv.WINDOW_NORMAL | cv.WINDOW_KEEPRATIO) cv.resizeWindow(display_name, 960, 720) cv.setMouseCallback(display_name, ui_control.mouse_callback) if hasattr(self, 'initialize_features'): self.initialize_features() while True: # Capture frame-by-frame ret, frame = cap.read() frame_disp = frame.copy() if ui_control.mode == 'track' and ui_control.mode_switch: ui_control.mode_switch = False init_state = ui_control.get_bb() self.initialize(frame, {'init_bbox': init_state}) # Draw box if ui_control.mode == 'select': cv.rectangle(frame_disp, ui_control.get_tl(), ui_control.get_br(), (255, 0, 0), 2) elif ui_control.mode == 'track': out = self.track(frame) state = [int(s) for s in out['target_bbox']] cv.rectangle(frame_disp, (state[0], state[1]), (state[2] + state[0], state[3] + state[1]), (0, 255, 0), 5) # Put text font_color = (0, 0, 0) if ui_control.mode == 'init' or ui_control.mode == 'select': cv.putText(frame_disp, 'Select target', (20, 30), cv.FONT_HERSHEY_COMPLEX_SMALL, 1, font_color, 1) cv.putText(frame_disp, 'Press q to quit', (20, 55), cv.FONT_HERSHEY_COMPLEX_SMALL, 1, font_color, 1) elif ui_control.mode == 'track': cv.putText(frame_disp, 'Tracking!', (20, 30), cv.FONT_HERSHEY_COMPLEX_SMALL, 1, font_color, 1) cv.putText(frame_disp, 'Press r to reset', (20, 55), cv.FONT_HERSHEY_COMPLEX_SMALL, 1, font_color, 1) cv.putText(frame_disp, 'Press q to quit', (20, 80), cv.FONT_HERSHEY_COMPLEX_SMALL, 1, font_color, 1) # Display the resulting frame cv.imshow(display_name, frame_disp) key = cv.waitKey(1) if key == ord('q'): break elif key == ord('r'): ui_control.mode = 'init' # When everything done, release the capture cap.release() cv.destroyAllWindows() def track_vot(self): """Run tracker on VOT.""" def _convert_anno_to_list(vot_anno): vot_anno = [ vot_anno[0][0][0], vot_anno[0][0][1], vot_anno[0][1][0], vot_anno[0][1][1], vot_anno[0][2][0], vot_anno[0][2][1], vot_anno[0][3][0], vot_anno[0][3][1] ] return vot_anno def _convert_image_path(image_path): image_path_new = image_path[20:-2] return "".join(image_path_new) handle = vot.VOT("polygon") vot_anno_polygon = handle.region() vot_anno_polygon = _convert_anno_to_list(vot_anno_polygon) init_state = convert_vot_anno_to_rect( vot_anno_polygon, self.params.vot_anno_conversion_type) image_path = handle.frame() if not image_path: return image_path = _convert_image_path(image_path) image = self._read_image(image_path) self.initialize(image, {'init_bbox': init_state}) if self.visdom is not None: self.visdom.register((image, init_state), 'Tracking', 1, 'Tracking') # Track while True: while True: if not self.pause_mode: break elif self.step: self.step = False break else: time.sleep(0.1) image_path = handle.frame() if not image_path: break image_path = _convert_image_path(image_path) image = self._read_image(image_path) out = self.track(image) state = out['target_bbox'] if self.visdom is not None: self.visdom.register((image, state), 'Tracking', 1, 'Tracking') handle.report(vot.Rectangle(state[0], state[1], state[2], state[3])) def reset_tracker(self): pass def press(self, event): if event.key == 'p': self.pause_mode = not self.pause_mode print("Switching pause mode!") elif event.key == 'r': self.reset_tracker() print("Resetting target pos to gt!") def init_visualization(self): # plt.ion() self.pause_mode = False self.fig, self.ax = plt.subplots(1) self.fig.canvas.mpl_connect('key_press_event', self.press) plt.tight_layout() def visualize(self, image, state, seq_name, frame_no): self.ax.cla() self.ax.imshow(image) rect = patches.Rectangle((state[0], state[1]), state[2], state[3], linewidth=1, edgecolor='r', facecolor='none') self.ax.add_patch(rect) if hasattr(self, 'gt_state') and False: gt_state = self.gt_state rect = patches.Rectangle((gt_state[0], gt_state[1]), gt_state[2], gt_state[3], linewidth=1, edgecolor='g', facecolor='none') self.ax.add_patch(rect) self.ax.set_axis_off() self.ax.axis('equal') #draw_figure(self.fig) if not isdir(seq_name): os.makedirs(seq_name) self.fig.savefig( join(seq_name, 'frame' + str(frame_no).zfill(5) + '.png')) if self.pause_mode: keypress = False while not keypress: keypress = plt.waitforbuttonpress() def show_image(self, im, plot_name=None, ax=None): if isinstance(im, torch.Tensor): im = torch_to_numpy(im) # plot_id = sum([ord(x) for x in list(plot_name)]) if ax is None: plot_fig_name = 'debug_fig_' + plot_name plot_ax_name = 'debug_ax_' + plot_name if not hasattr(self, plot_fig_name): fig, ax = plt.subplots(1) setattr(self, plot_fig_name, fig) setattr(self, plot_ax_name, ax) plt.tight_layout() ax.set_title(plot_name) else: fig = getattr(self, plot_fig_name, None) ax = getattr(self, plot_ax_name, None) ax.cla() ax.imshow(im) ax.set_axis_off() ax.axis('equal') ax.set_title(plot_name) draw_figure(fig) def _read_image(self, image_file: str): return cv.cvtColor(cv.imread(image_file), cv.COLOR_BGR2RGB)
class BaseTracker: """Base class for all trackers.""" def visdom_ui_handler(self, data): if data['event_type'] == 'KeyPress': if data['key'] == ' ': self.pause_mode = not self.pause_mode elif data['key'] == 'ArrowRight' and self.pause_mode: self.step = True def __init__(self, params): self.params = params self.pause_mode = False self.step = False self.visdom = None if self.params.debug > 0 and self.params.visdom_info.get( 'use_visdom', True): try: self.visdom = Visdom(self.params.debug, { 'handler': self.visdom_ui_handler, 'win_id': 'Tracking' }, visdom_info=self.params.visdom_info) # Show help help_text = 'You can pause/unpause the tracker by pressing ''space'' with the ''Tracking'' window ' \ 'selected. During paused mode, you can track for one frame by pressing the right arrow key.' \ 'To enable/disable plotting of a data block, tick/untick the corresponding entry in ' \ 'block list.' self.visdom.register(help_text, 'text', 1, 'Help') except: time.sleep(0.5) print( '!!! WARNING: Visdom could not start, so using matplotlib visualization instead !!!\n' '!!! Start Visdom in a separate terminal window by typing \'visdom\' !!!' ) def initialize(self, image, info: dict) -> dict: """Overload this function in your tracker. This should initialize the model.""" raise NotImplementedError def track(self, image) -> dict: """Overload this function in your tracker. This should track in the frame and update the model.""" raise NotImplementedError def track_sequence(self, sequence): """Run tracker on a sequence.""" output = {'target_bbox': [], 'time': [], 'scores': []} def _store_outputs(tracker_out: dict, defaults=None): defaults = {} if defaults is None else defaults for key in tracker_out.keys(): if key not in output: raise RuntimeError('Unknown output from tracker.') for key in output.keys(): val = tracker_out.get(key, defaults.get(key, None)) if val is not None: output[key].append(val) # Initialize image = self._read_image(sequence.frames[0]) #['image'] (480, 640, 3) if hasattr(self.params, 'use_depth_channel'): if self.params.use_depth_channel: print('have %d depth frames' % (len(sequence.depth_frames))) depth = self._read_depth(sequence.depth_frames[0]) #depth = depth/1000.0 #from mm to m depth = np.repeat(np.expand_dims(depth, axis=2), 3, axis=2) #.astype(np.uint8) if image.shape[0] != depth.shape[0]: depth = depth[1:image.shape[0] + 1, :, :] #image = np.concatenate((image,np.expand_dims(depth,axis=2)),axis=2).astype(np.uint8) #print(['image'],image.shape) #['image'] (480, 640, 4) #print(['depth', depth.shape, np.mean(depth, (0,1)), np.std(depth, (0,1))]) #['depth', (480, 640), 22.48,19.60] if self.params.visualization and self.visdom is None: self.init_visualization() self.visualize(image[:, :, 0:3], sequence.get('init_bbox')) start_time = time.time() if hasattr(self.params, 'use_depth_channel'): out = self.initialize(image, depth, sequence.init_info()) else: out = self.initialize(image, sequence.init_info()) if out is None: out = {} _store_outputs( out, { 'target_bbox': sequence.get('init_bbox'), 'time': time.time() - start_time, 'scores': 1.0 }) if self.visdom is not None: self.visdom.register((image, sequence.get('init_bbox')), 'Tracking', 1, 'Tracking') # Track ind_frame = 0 for frame in sequence.frames[1:]: ind_frame = ind_frame + 1 self.ind_frame = ind_frame while True: if not self.pause_mode: break elif self.step: self.step = False break else: time.sleep(0.1) image = self._read_image(frame) if hasattr(self.params, 'use_depth_channel'): #print(['depth image',sequence.depth_frames[ind_frame]]) depth = self._read_depth(sequence.depth_frames[ind_frame]) #depth = depth/1000.0 #from mm to m depth = np.repeat(np.expand_dims(depth, axis=2), 3, axis=2) #.astype(np.uint8) start_time = time.time() if hasattr(self.params, 'use_depth_channel'): out = self.track(image, depth) else: out = self.track(image) _store_outputs( out, { 'time': time.time() - start_time, 'scores': self.debug_info['max_score'] }) #get gt_state if the gt_state for the whole sequence is provided if sequence.ground_truth_rect.shape[0] > 1: self.gt_state = sequence.ground_truth_rect[ind_frame] if self.visdom is not None: self.visdom.register((image, out['target_bbox']), 'Tracking', 1, 'Tracking') elif self.params.visualization: # if hasattr(self.params, 'use_depth_channel'): # self.visualize(image, out['target_bbox'], out_rgb['target_bbox'], out_depth['target_bbox']) # else: self.visualize(image, out['target_bbox']) #visualize the depth if hasattr(self.params, 'use_depth_channel'): if os.path.exists(sequence.depth_frames[ind_frame]): #dimage=self._read_image(sequence.depth_frames[ind_frame]) self.visualize_depth( np.uint8(255 * depth / np.max(depth)), out['target_bbox']) #print(depth.shape) pass return output def track_videofile(self, videofilepath, optional_box=None): """Run track with a video file input.""" assert os.path.isfile(videofilepath), "Invalid param {}".format( videofilepath) ", videofilepath must be a valid videofile" if hasattr(self, 'initialize_features'): self.initialize_features() cap = cv.VideoCapture(videofilepath) display_name = 'Display: ' + self.params.tracker_name cv.namedWindow(display_name, cv.WINDOW_NORMAL | cv.WINDOW_KEEPRATIO) cv.resizeWindow(display_name, 960, 720) success, frame = cap.read() cv.imshow(display_name, frame) if success is not True: print("Read frame from {} failed.".format(videofilepath)) exit(-1) if optional_box is not None: assert isinstance(optional_box, list, tuple) assert len(optional_box) == 4, "valid box's foramt is [x,y,w,h]" self.initialize(frame, {'init_bbox': optional_box}) else: while True: # cv.waitKey() frame_disp = frame.copy() cv.putText(frame_disp, 'Select target ROI and press ENTER', (20, 30), cv.FONT_HERSHEY_COMPLEX_SMALL, 1.5, (0, 0, 0), 1) x, y, w, h = cv.selectROI(display_name, frame_disp, fromCenter=False) init_state = [x, y, w, h] self.initialize(frame, {'init_bbox': init_state}) break while True: ret, frame = cap.read() if frame is None: return frame_disp = frame.copy() # Draw box out = self.track(frame) state = [int(s) for s in out['target_bbox']] cv.rectangle(frame_disp, (state[0], state[1]), (state[2] + state[0], state[3] + state[1]), (0, 255, 0), 5) font_color = (0, 0, 0) cv.putText(frame_disp, 'Tracking!', (20, 30), cv.FONT_HERSHEY_COMPLEX_SMALL, 1, font_color, 1) cv.putText(frame_disp, 'Press r to reset', (20, 55), cv.FONT_HERSHEY_COMPLEX_SMALL, 1, font_color, 1) cv.putText(frame_disp, 'Press q to quit', (20, 80), cv.FONT_HERSHEY_COMPLEX_SMALL, 1, font_color, 1) # Display the resulting frame cv.imshow(display_name, frame_disp) key = cv.waitKey(1) if key == ord('q'): break elif key == ord('r'): ret, frame = cap.read() frame_disp = frame.copy() cv.putText(frame_disp, 'Select target ROI and press ENTER', (20, 30), cv.FONT_HERSHEY_COMPLEX_SMALL, 1.5, (0, 0, 0), 1) cv.imshow(display_name, frame_disp) x, y, w, h = cv.selectROI(display_name, frame_disp, fromCenter=False) init_state = [x, y, w, h] self.initialize(frame, {'init_bbox': init_state}) # When everything done, release the capture cap.release() cv.destroyAllWindows() def track_webcam(self): """Run tracker with webcam.""" class UIControl: def __init__(self): self.mode = 'init' # init, select, track self.target_tl = (-1, -1) self.target_br = (-1, -1) self.mode_switch = False def mouse_callback(self, event, x, y, flags, param): if event == cv.EVENT_LBUTTONDOWN and self.mode == 'init': self.target_tl = (x, y) self.target_br = (x, y) self.mode = 'select' self.mode_switch = True elif event == cv.EVENT_MOUSEMOVE and self.mode == 'select': self.target_br = (x, y) elif event == cv.EVENT_LBUTTONDOWN and self.mode == 'select': self.target_br = (x, y) self.mode = 'track' self.mode_switch = True def get_tl(self): return self.target_tl if self.target_tl[0] < self.target_br[ 0] else self.target_br def get_br(self): return self.target_br if self.target_tl[0] < self.target_br[ 0] else self.target_tl def get_bb(self): tl = self.get_tl() br = self.get_br() bb = [ min(tl[0], br[0]), min(tl[1], br[1]), abs(br[0] - tl[0]), abs(br[1] - tl[1]) ] return bb ui_control = UIControl() cap = cv.VideoCapture(0) display_name = 'Display: ' + self.params.tracker_name cv.namedWindow(display_name, cv.WINDOW_NORMAL | cv.WINDOW_KEEPRATIO) cv.resizeWindow(display_name, 960, 720) cv.setMouseCallback(display_name, ui_control.mouse_callback) if hasattr(self, 'initialize_features'): self.initialize_features() while True: # Capture frame-by-frame ret, frame = cap.read() frame_disp = frame.copy() if ui_control.mode == 'track' and ui_control.mode_switch: ui_control.mode_switch = False init_state = ui_control.get_bb() self.initialize(frame, {'init_bbox': init_state}) # Draw box if ui_control.mode == 'select': cv.rectangle(frame_disp, ui_control.get_tl(), ui_control.get_br(), (255, 0, 0), 2) elif ui_control.mode == 'track': out = self.track(frame) state = [int(s) for s in out['target_bbox']] cv.rectangle(frame_disp, (state[0], state[1]), (state[2] + state[0], state[3] + state[1]), (0, 255, 0), 5) # Put text font_color = (0, 0, 0) if ui_control.mode == 'init' or ui_control.mode == 'select': cv.putText(frame_disp, 'Select target', (20, 30), cv.FONT_HERSHEY_COMPLEX_SMALL, 1, font_color, 1) cv.putText(frame_disp, 'Press q to quit', (20, 55), cv.FONT_HERSHEY_COMPLEX_SMALL, 1, font_color, 1) elif ui_control.mode == 'track': cv.putText(frame_disp, 'Tracking!', (20, 30), cv.FONT_HERSHEY_COMPLEX_SMALL, 1, font_color, 1) cv.putText(frame_disp, 'Press r to reset', (20, 55), cv.FONT_HERSHEY_COMPLEX_SMALL, 1, font_color, 1) cv.putText(frame_disp, 'Press q to quit', (20, 80), cv.FONT_HERSHEY_COMPLEX_SMALL, 1, font_color, 1) # Display the resulting frame cv.imshow(display_name, frame_disp) key = cv.waitKey(1) if key == ord('q'): break elif key == ord('r'): ui_control.mode = 'init' # When everything done, release the capture cap.release() cv.destroyAllWindows() def track_vot(self): """Run tracker on VOT.""" def _convert_anno_to_list(vot_anno): vot_anno = [ vot_anno[0][0][0], vot_anno[0][0][1], vot_anno[0][1][0], vot_anno[0][1][1], vot_anno[0][2][0], vot_anno[0][2][1], vot_anno[0][3][0], vot_anno[0][3][1] ] return vot_anno def _convert_image_path(image_path): image_path_new = image_path[20:-2] return "".join(image_path_new) handle = vot.VOT("polygon") vot_anno_polygon = handle.region() vot_anno_polygon = _convert_anno_to_list(vot_anno_polygon) init_state = convert_vot_anno_to_rect( vot_anno_polygon, self.params.vot_anno_conversion_type) image_path = handle.frame() if not image_path: return image_path = _convert_image_path(image_path) image = self._read_image(image_path) self.initialize(image, {'init_bbox': init_state}) if self.visdom is not None: self.visdom.register((image, init_state), 'Tracking', 1, 'Tracking') # Track while True: while True: if not self.pause_mode: break elif self.step: self.step = False break else: time.sleep(0.1) image_path = handle.frame() if not image_path: break image_path = _convert_image_path(image_path) image = self._read_image(image_path) out = self.track(image) state = out['target_bbox'] if self.visdom is not None: self.visdom.register((image, state), 'Tracking', 1, 'Tracking') handle.report(vot.Rectangle(state[0], state[1], state[2], state[3])) def reset_tracker(self): pass def press(self, event): if event.key == 'p': self.pause_mode = not self.pause_mode print("Switching pause mode!") elif event.key == 'r': self.reset_tracker() print("Resetting target pos to gt!") def init_visualization(self): # plt.ion() self.pause_mode = False self.fig, self.ax = plt.subplots(1) self.fig2, self.ax2 = plt.subplots(1) self.fig.canvas.mpl_connect('key_press_event', self.press) plt.tight_layout() def visualize(self, image, state, *var): self.ax.cla() self.ax.imshow(image) if (state[2] != 0 and state[3] != 0): self.ax.text(10, 30, 'FOUND', fontsize=14, bbox=dict(facecolor='green', alpha=0.2)) else: self.ax.text(10, 30, 'NOT FOUND', fontsize=14, bbox=dict(facecolor='red', alpha=0.2)) pass if len(var) == 0: rect = patches.Rectangle((state[0], state[1]), state[2], state[3], linewidth=1, edgecolor='r', facecolor='none') self.ax.add_patch(rect) if len(var) > 0: #state_rgb, state_depth, provided state_rgb = var[0] state_depth = var[1] #draw one dot for the center of state_rgb # self.ax.plot(state_rgb[0]+state_rgb[2]/2, state_rgb[1]+state_rgb[3]/2,'ro') # self.ax.plot(state_depth[0]+state_depth[2]/2, state_depth[1]+state_depth[3]/2, 'bo') # self.ax.plot(state[0]+state[2]/2, state[1]+state[3]/2,'wo') #another dot for the center of state_depth # rect_rgb= patches.Rectangle((state_rgb[0], state_rgb[1]), state_rgb[2], state_rgb[3], linewidth=2, edgecolor='r', facecolor='none') # self.ax.add_patch(rect_rgb) rect_depth = patches.Rectangle((state_depth[0], state_depth[1]), state_depth[2], state_depth[3], linewidth=2, edgecolor='b', facecolor='none') self.ax.add_patch(rect_depth) rect = patches.Rectangle((state[0], state[1]), state[2], state[3], linewidth=1, edgecolor='w', facecolor='none') self.ax.add_patch(rect) #print(['var', var]) #['var', (tensor([263.5000, 266.5000]), tensor([263.5000, 266.5000]), tensor([263.6045, 271.1568]))] if hasattr(self, 'gt_state') and True: gt_state = self.gt_state self.ax.plot(gt_state[0] + gt_state[2] / 2, gt_state[1] + gt_state[3] / 2, 'go') rect = patches.Rectangle((gt_state[0], gt_state[1]), gt_state[2], gt_state[3], linewidth=2, edgecolor='g', facecolor='none') self.ax.add_patch(rect) self.ax.set_axis_off() self.ax.axis('equal') draw_figure(self.fig) if hasattr(self, 'ind_frame'): if os.path.exists('./tracking_results/imgs'): self.fig.savefig('./tracking_results/imgs/img_%d.png' % self.ind_frame) if self.pause_mode: keypress = False while not keypress: keypress = plt.waitforbuttonpress() def visualize_depth(self, image, state): self.ax2.cla() self.ax2.imshow(image) self.ax2.set_axis_off() self.ax2.axis('equal') plt.draw() plt.pause(0.001) if hasattr(self, 'ind_frame'): if os.path.exists('./tracking_results/imgs'): self.fig2.savefig('./tracking_results/imgs/depth_%d.png' % self.ind_frame) if self.pause_mode: plt.waitforbuttonpress() def show_image(self, im, plot_name=None, ax=None): if isinstance(im, torch.Tensor): im = torch_to_numpy(im) # plot_id = sum([ord(x) for x in list(plot_name)]) if ax is None: plot_fig_name = 'debug_fig_' + plot_name plot_ax_name = 'debug_ax_' + plot_name if not hasattr(self, plot_fig_name): fig, ax = plt.subplots(1) setattr(self, plot_fig_name, fig) setattr(self, plot_ax_name, ax) plt.tight_layout() ax.set_title(plot_name) else: fig = getattr(self, plot_fig_name, None) ax = getattr(self, plot_ax_name, None) ax.cla() ax.imshow(im) ax.set_axis_off() ax.axis('equal') ax.set_title(plot_name) draw_figure(fig) def _read_depth(self, image_file: str): # Full kernels FULL_KERNEL_3 = np.ones((3, 3), np.uint8) FULL_KERNEL_5 = np.ones((5, 5), np.uint8) FULL_KERNEL_7 = np.ones((7, 7), np.uint8) FULL_KERNEL_9 = np.ones((9, 9), np.uint8) FULL_KERNEL_31 = np.ones((31, 31), np.uint8) depth = cv.imread(image_file, cv.COLOR_BGR2GRAY) #print(['_read_depth', depth.min(), depth.max(), depth.mean(), depth.std()]) if 'Princeton' in image_file: #depth.max()>=60000: # bug found, we need to bitshift depth. depth = np.bitwise_or(np.right_shift(depth, 3), np.left_shift(depth, 13)) depth = depth / 1000.0 depth[depth >= 8.0] = 8.0 depth[depth <= 0.0] = 8.0 #depth=8.0-depth # Hole closing depth = cv.morphologyEx(depth, cv.MORPH_CLOSE, FULL_KERNEL_7) #depth = 255.0*depth/(np.max(depth)+1e-3) return depth def _read_image(self, image_file: str): return cv.cvtColor(cv.imread(image_file), cv.COLOR_BGR2RGB)
class TrackerYolo: """Wraps the tracker for evaluation and running purposes. args: name: Name of tracking method. parameter_name: Name of parameter file. run_id: The run id. display_name: Name to be displayed in the result plots. """ def __init__(self, name: str, parameter_name: str, run_id: int = None, display_name: str = None): assert run_id is None or isinstance(run_id, int) self.name = name self.parameter_name = parameter_name self.run_id = run_id self.display_name = display_name env = env_settings() if self.run_id is None: self.results_dir = '{}/{}/{}'.format(env.results_path, self.name, self.parameter_name) self.segmentation_dir = '{}/{}/{}'.format(env.segmentation_path, self.name, self.parameter_name) else: self.results_dir = '{}/{}/{}_{:03d}'.format(env.results_path, self.name, self.parameter_name, self.run_id) self.segmentation_dir = '{}/{}/{}_{:03d}'.format(env.segmentation_path, self.name, self.parameter_name, self.run_id) tracker_module_abspath = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', 'tracker', self.name)) if os.path.isdir(tracker_module_abspath): tracker_module = importlib.import_module('pytracking.tracker.{}'.format(self.name)) self.tracker_class = tracker_module.get_tracker_class() else: self.tracker_class = None self.visdom = None def _init_visdom(self, visdom_info, debug): visdom_info = {} if visdom_info is None else visdom_info self.pause_mode = False self.step = False if debug > 0 and visdom_info.get('use_visdom', True): try: self.visdom = Visdom(debug, {'handler': self._visdom_ui_handler, 'win_id': 'Tracking'}, visdom_info=visdom_info) # Show help help_text = 'You can pause/unpause the tracker by pressing ''space'' with the ''Tracking'' window ' \ 'selected. During paused mode, you can track for one frame by pressing the right arrow key.' \ 'To enable/disable plotting of a data block, tick/untick the corresponding entry in ' \ 'block list.' self.visdom.register(help_text, 'text', 1, 'Help') except: time.sleep(0.5) print('!!! WARNING: Visdom could not start, so using matplotlib visualization instead !!!\n' '!!! Start Visdom in a separate terminal window by typing \'visdom\' !!!') def _visdom_ui_handler(self, data): if data['event_type'] == 'KeyPress': if data['key'] == ' ': self.pause_mode = not self.pause_mode elif data['key'] == 'ArrowRight' and self.pause_mode: self.step = True def create_tracker(self, params): tracker = self.tracker_class(params) tracker.visdom = self.visdom return tracker def run_video(self, videofilepath, optional_box=None, debug=None, visdom_info=None, save_results=False): """Run the tracker with the vieofile. args: debug: Debug level. """ def yolo_search(W, H, frame_yolo): # if the frame dimensions are empty, grab them if W is None or H is None: (H, W) = frame_yolo.shape[:2] # construct a blob from the input frame and then perform a forward # pass of the YOLO object detector, giving us our bounding boxes # and associated probabilities blob = cv.dnn.blobFromImage(frame_yolo, 1 / 255.0, (416, 416), swapRB=True, crop=False) net.setInput(blob) layerOutputs = net.forward(ln) # initialize our lists of detected bounding boxes, confidences, # and class IDs, respectively boxes = [] confidences = [] classIDs = [] # loop over each of the layer outputs for output in layerOutputs: # loop over each of the detections for detection in output: # extract the class ID and confidence (i.e., probability) # of the current object detection scores = detection[5:] classID = np.argmax(scores) confidence = scores[classID] # filter weak prediction and unrelated classes if classID not in outdoor_classes and confidence > 0.5: # scale the bounding box coordinates back relative to # the size of the image, keeping in mind that YOLO # actually returns the center (x, y)-coordinates of # the bounding box followed by the boxes' width and # height box = detection[0:4] * np.array([W, H, W, H]) (centerX, centerY, width, height) = box.astype("int") # use the center (x, y)-coordinates to derive the top # and and left corner of the bounding box x = int(centerX - (width / 2)) y = int(centerY - (height / 2)) # update our list of bounding box coordinates, # confidences, and class IDs boxes.append([x, y, int(width), int(height)]) confidences.append(float(confidence)) classIDs.append(classID) # apply non-maxima suppression to suppress weak, overlapping # bounding boxes idxs = cv.dnn.NMSBoxes(boxes, confidences, 0.5, 0.3) # ensure at least one detection exists if len(idxs) > 0: # loop over the indexes we are keeping for i in idxs.flatten(): # extract the bounding box coordinates (x, y) = (boxes[i][0], boxes[i][1]) (w, h) = (boxes[i][2], boxes[i][3]) # draw a bounding box rectangle and label on the frame color = [int(c) for c in COLORS[classIDs[i]]] cv.rectangle(frame_yolo, (x, y), (x + w, y + h), color, 2) text = "{}: {:.4f}".format(LABELS[classIDs[i]], confidences[i]) cv.putText(frame_yolo, text, (x, y - 5), cv.FONT_HERSHEY_SIMPLEX, 0.5, color, 2) if classIDs[i] == 45: # 0 - person, 65 - remote 45 - bowl detection_flag = 1 tl_coor = (x, y) # top left coordinates br_coor = ((x + w), (y + h)) # bottom right coordinates cv.rectangle(frame_yolo, tl_coor, br_coor, (255, 255, 255), 2) return tl_coor, br_coor, detection_flag, frame_yolo return (0, 0), (0, 0), 0, frame_yolo # load the COCO class labels our YOLO model was trained on # and the classes that wont be used (coco.names contains the names) # Init a detection flag det_flag = 0 stop_yolo = 0 labelsPath = "/home/ebo/Parrot_CV_Project/local_yolo/yolo-coco/coco.names" LABELS = open(labelsPath).read().strip().split("\n") outdoor_classes = [2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 17, 18, 19, 20, 21, 22, 23, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38] # initialize a list of colors to represent each possible class label np.random.seed(42) COLORS = np.random.randint(0, 255, size=(len(LABELS), 3), dtype="uint8") # derive the paths to the YOLO weights and model configuration weightsPath = "/home/ebo/Parrot_CV_Project/local_yolo/yolo-coco/yolov3.weights" configPath = "/home/ebo/Parrot_CV_Project/local_yolo/yolo-coco/yolov3.cfg" # load our YOLO object detector trained on COCO dataset (80 classes) print("[INFO] loading YOLO from disk...") net = cv.dnn.readNetFromDarknet(configPath, weightsPath) # determine only the output layers from yolo ln = net.getLayerNames() ln = [ln[i[0] - 1] for i in net.getUnconnectedOutLayers()] (W, H) = (None, None) params = self.get_parameters() debug_ = debug if debug is None: debug_ = getattr(params, 'debug', 0) params.debug = debug_ params.tracker_name = self.name params.param_name = self.parameter_name self._init_visdom(visdom_info, debug_) multiobj_mode = getattr(params, 'multiobj_mode', getattr(self.tracker_class, 'multiobj_mode', 'default')) if multiobj_mode == 'default': tracker = self.create_tracker(params) if hasattr(tracker, 'initialize_features'): tracker.initialize_features() elif multiobj_mode == 'parallel': tracker = MultiObjectWrapper(self.tracker_class, params, self.visdom, fast_load=True) else: raise ValueError('Unknown multi object mode {}'.format(multiobj_mode)) assert os.path.isfile(videofilepath), "Invalid param {}".format(videofilepath) ", videofilepath must be a valid videofile" output_boxes = [] cap = cv.VideoCapture(videofilepath) display_name = 'Display: ' + tracker.params.tracker_name cv.namedWindow(display_name, cv.WINDOW_NORMAL | cv.WINDOW_KEEPRATIO) cv.resizeWindow(display_name, 960, 720) success, frame = cap.read() cv.imshow(display_name, frame) def _build_init_info(box): return {'init_bbox': OrderedDict({1: box}), 'init_object_ids': [1, ], 'object_ids': [1, ], 'sequence_object_ids': [1, ]} if success is not True: print("Read frame from {} failed.".format(videofilepath)) exit(-1) while True: ret, frame = cap.read() if frame is None: break frame_disp = frame.copy() if W is None or H is None: (H, W) = frame_disp.shape[:2] if stop_yolo == 0: tl_yolo, br_yolo, det_flag, frame_disp = yolo_search(W, H, frame.copy()) cv.putText(frame_disp, "Searching: BOWL", (50, 50), cv.FONT_HERSHEY_SIMPLEX, 1, (255, 0, 0), 2) if det_flag == 1: if stop_yolo == 0: stop_yolo = 1 x = tl_yolo[0] y = tl_yolo[1] w = abs(br_yolo[0] - tl_yolo[0]) h = abs(br_yolo[1] - tl_yolo[1]) init_state = [x, y, w, h] tracker.initialize(frame, _build_init_info(init_state)) output_boxes.append(init_state) # Draw box out = tracker.track(frame) state = [int(s) for s in out['target_bbox'][1]] output_boxes.append(state) tl = (state[0], state[1]) br = (state[2] + state[0], state[3] + state[1]) w = state[2] h = state[3] cv.rectangle(frame_disp, tl, br, (0, 255, 0), 5) center = (int(tl[0] + w/2), int(tl[1] + h/2)) cv.circle(frame_disp, center, 3, (0, 0, 255), -1) cv.putText(frame_disp, "FOUND BOWL", (50, 50), cv.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 2) if center[0] < W*0.40: cv.putText(frame_disp, "MOVE LEFT", (50, 150), cv.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 2) elif center[0] > W*0.60: cv.putText(frame_disp, "MOVE RIGHT", (450, 150), cv.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 2) if center[1] < H*0.40: cv.putText(frame_disp, "MOVE UP", (200, 50), cv.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 2) elif center[1] > H*0.60: cv.putText(frame_disp, "MOVE DOWN", (200, 300), cv.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 2) if w*h < W*H*0.05: cv.putText(frame_disp, "MOVE FORWARD", (int(W/2), int(H/2)), cv.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 2) elif w*h > W*H*0.15: cv.putText(frame_disp, "MOVE BACK", (int(W/2), int(H/2)), cv.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 2) # Display the resulting frame cv.imshow(display_name, frame_disp) key = cv.waitKey(1) if key == ord('q'): break # When everything done, release the capture cap.release() cv.destroyAllWindows() if save_results: if not os.path.exists(self.results_dir): os.makedirs(self.results_dir) video_name = Path(videofilepath).stem base_results_path = os.path.join(self.results_dir, 'video_{}'.format(video_name)) tracked_bb = np.array(output_boxes).astype(int) bbox_file = '{}.txt'.format(base_results_path) np.savetxt(bbox_file, tracked_bb, delimiter='\t', fmt='%d') def run_webcam(self, debug=None, visdom_info=None): """Run the tracker with the webcam. args: debug: Debug level. """ def yolo_search(W, H, frame_yolo): fl = 0 # if the frame dimensions are empty, grab them if W is None or H is None: (H, W) = frame_yolo.shape[:2] # construct a blob from the input frame and then perform a forward # pass of the YOLO object detector, giving us our bounding boxes # and associated probabilities blob = cv.dnn.blobFromImage(frame_yolo, 1 / 255.0, (416, 416), swapRB=True, crop=False) net.setInput(blob) layerOutputs = net.forward(ln) # initialize our lists of detected bounding boxes, confidences, # and class IDs, respectively boxes = [] confidences = [] classIDs = [] # loop over each of the layer outputs for output in layerOutputs: # loop over each of the detections for detection in output: # extract the class ID and confidence (i.e., probability) # of the current object detection scores = detection[5:] classID = np.argmax(scores) confidence = scores[classID] # filter weak prediction and unrelated classes if classID not in outdoor_classes and confidence > 0.5: # scale the bounding box coordinates back relative to # the size of the image, keeping in mind that YOLO # actually returns the center (x, y)-coordinates of # the bounding box followed by the boxes' width and # height box = detection[0:4] * np.array([W, H, W, H]) (centerX, centerY, width, height) = box.astype("int") # use the center (x, y)-coordinates to derive the top # and and left corner of the bounding box x = int(centerX - (width / 2)) y = int(centerY - (height / 2)) # update our list of bounding box coordinates, # confidences, and class IDs boxes.append([x, y, int(width), int(height)]) confidences.append(float(confidence)) classIDs.append(classID) # apply non-maxima suppression to suppress weak, overlapping # bounding boxes idxs = cv.dnn.NMSBoxes(boxes, confidences, 0.5, 0.3) # ensure at least one detection exists if len(idxs) > 0: # loop over the indexes we are keeping for i in idxs.flatten(): # extract the bounding box coordinates (x, y) = (boxes[i][0], boxes[i][1]) (w, h) = (boxes[i][2], boxes[i][3]) # draw a bounding box rectangle and label on the frame color = [int(c) for c in COLORS[classIDs[i]]] cv.rectangle(frame_yolo, (x, y), (x + w, y + h), color, 2) text = "{}: {:.4f}".format(LABELS[classIDs[i]], confidences[i]) cv.putText(frame_yolo, text, (x, y - 5), cv.FONT_HERSHEY_SIMPLEX, 0.5, color, 2) if classIDs[i] == 0: detection_flag = 1 tl_coor = (x, y) # top left coordinates br_coor = ((x + w), (y + h)) # bottom right coordinates coordinates_text = "{} {}".format(tl_coor, br_coor) cv.rectangle(frame_yolo, tl_coor, br_coor, (255, 255, 255), 2) fl = 1 return tl_coor, br_coor, detection_flag, frame_yolo if fl == 0: return (0, 0), (0, 0), 0, frame_yolo # load the COCO class labels our YOLO model was trained on # and the classes that wont be used (coco.names contains the names) # Init a detection flag det_flag = 0 labelsPath = "/home/ebo/Parrot_CV_Project/local_yolo/yolo-coco/coco.names" LABELS = open(labelsPath).read().strip().split("\n") outdoor_classes = [2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 17, 18, 19, 20, 21, 22, 23, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38] # initialize a list of colors to represent each possible class label np.random.seed(42) COLORS = np.random.randint(0, 255, size=(len(LABELS), 3), dtype="uint8") # derive the paths to the YOLO weights and model configuration weightsPath = "/home/ebo/Parrot_CV_Project/local_yolo/yolo-coco/yolov3.weights" configPath = "/home/ebo/Parrot_CV_Project/local_yolo/yolo-coco/yolov3.cfg" # load our YOLO object detector trained on COCO dataset (80 classes) print("[INFO] loading YOLO from disk...") net = cv.dnn.readNetFromDarknet(configPath, weightsPath) # determine only the output layers from yolo ln = net.getLayerNames() ln = [ln[i[0] - 1] for i in net.getUnconnectedOutLayers()] (W, H) = (None, None) temp_flag = 0 params = self.get_parameters() debug_ = debug if debug is None: debug_ = getattr(params, 'debug', 0) params.debug = debug_ params.tracker_name = self.name params.param_name = self.parameter_name self._init_visdom(visdom_info, debug_) multiobj_mode = getattr(params, 'multiobj_mode', getattr(self.tracker_class, 'multiobj_mode', 'default')) if multiobj_mode == 'default': tracker = self.create_tracker(params) elif multiobj_mode == 'parallel': tracker = MultiObjectWrapper(self.tracker_class, params, self.visdom, fast_load=True) else: raise ValueError('Unknown multi object mode {}'.format(multiobj_mode)) class UIControl: def __init__(self): self.mode = 'init' # init, select, track self.new_init = False def get_bb(self): # yolo bb if det_flag == 1: tl = tl_yolo br = br_yolo bb = [min(tl[0], br[0]), min(tl[1], br[1]), abs(br[0] - tl[0]), abs(br[1] - tl[1])] return bb ui_control = UIControl() cap = cv.VideoCapture(0) display_name = 'Display: ' + self.name cv.namedWindow(display_name, cv.WINDOW_NORMAL | cv.WINDOW_KEEPRATIO) cv.resizeWindow(display_name, 960, 720) next_object_id = 1 sequence_object_ids = [] prev_output = OrderedDict() while True: # Capture frame-by-frame ret, frame = cap.read() frame_disp = frame.copy() tl_yolo, br_yolo, det_flag, frame_yolo = yolo_search(W, H, frame.copy()) info = OrderedDict() info['previous_output'] = prev_output # If there's a human detection, show it if det_flag == 1 and temp_flag == 0: init_state = ui_control.get_bb() info['init_object_ids'] = [next_object_id, ] info['init_bbox'] = OrderedDict({next_object_id: init_state}) sequence_object_ids.append(next_object_id) next_object_id += 1 temp_flag = 1 if len(sequence_object_ids) > 0: info['sequence_object_ids'] = sequence_object_ids out = tracker.track(frame, info) prev_output = OrderedDict(out) if 'segmentation' in out: frame_disp = overlay_mask(frame_disp, out['segmentation']) if 'target_bbox' in out: for obj_id, state in out['target_bbox'].items(): state = [int(s) for s in state] cv.rectangle(frame_disp, (state[0], state[1]), (state[2] + state[0], state[3] + state[1]), _tracker_disp_colors[obj_id], 5) # Put text font_color = (0, 0, 0) cv.putText(frame_disp, 'Press r to reset', (20, 25), cv.FONT_HERSHEY_COMPLEX_SMALL, 1, font_color, 1) cv.putText(frame_disp, 'Press q to quit', (20, 55), cv.FONT_HERSHEY_COMPLEX_SMALL, 1, font_color, 1) # Display the resulting frame cv.imshow(display_name, frame_disp) cv.imshow("YOLO", frame_yolo) key = cv.waitKey(1) if key == ord('q'): break elif key == ord('r'): next_object_id = 1 sequence_object_ids = [] prev_output = OrderedDict() info = OrderedDict() info['object_ids'] = [] info['init_object_ids'] = [] info['init_bbox'] = OrderedDict() tracker.initialize(frame, info) ui_control.mode = 'init' # When everything done, release the capture cap.release() cv.destroyAllWindows() def get_parameters(self): """Get parameters.""" param_module = importlib.import_module('pytracking.parameter.{}.{}'.format(self.name, self.parameter_name)) params = param_module.parameters() return params def init_visualization(self): self.pause_mode = False self.fig, self.ax = plt.subplots(1) self.fig.canvas.mpl_connect('key_press_event', self.press) plt.tight_layout() def visualize(self, image, state, segmentation=None): self.ax.cla() self.ax.imshow(image) if segmentation is not None: self.ax.imshow(segmentation, alpha=0.5) if isinstance(state, (OrderedDict, dict)): boxes = [v for k, v in state.items()] else: boxes = (state,) for i, box in enumerate(boxes, start=1): col = _tracker_disp_colors[i] col = [float(c) / 255.0 for c in col] rect = patches.Rectangle((box[0], box[1]), box[2], box[3], linewidth=1, edgecolor=col, facecolor='none') self.ax.add_patch(rect) if getattr(self, 'gt_state', None) is not None: gt_state = self.gt_state rect = patches.Rectangle((gt_state[0], gt_state[1]), gt_state[2], gt_state[3], linewidth=1, edgecolor='g', facecolor='none') self.ax.add_patch(rect) self.ax.set_axis_off() self.ax.axis('equal') draw_figure(self.fig) if self.pause_mode: keypress = False while not keypress: keypress = plt.waitforbuttonpress() def reset_tracker(self): pass def press(self, event): if event.key == 'p': self.pause_mode = not self.pause_mode print("Switching pause mode!") elif event.key == 'r': self.reset_tracker() print("Resetting target pos to gt!") def _read_image(self, image_file: str): im = cv.imread(image_file) return cv.cvtColor(im, cv.COLOR_BGR2RGB)