Support video input

Author: MasterBin-IIAU

lib/test/evaluation/tracker.py

@@ -6,6 +6,8 @@
 import cv2 as cv
 
 from lib.utils.lmdb_utils import decode_img
+from pathlib import Path
+import numpy as np
 
 
 def trackerlist(name: str, parameter_name: str, dataset_name: str, run_ids = None, display_name: str = None,
@@ -147,6 +149,128 @@ def _store_outputs(tracker_out: dict, defaults=None):
 
         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)
+
+        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 get_parameters(self):
         """Get parameters."""
         param_module = importlib.import_module('lib.test.parameter.{}'.format(self.name))

tracking/video_demo.py

@@ -0,0 +1,39 @@
+import os
+import sys
+import argparse
+
+prj_path = os.path.join(os.path.dirname(__file__), '..')
+if prj_path not in sys.path:
+    sys.path.append(prj_path)
+
+from lib.test.evaluation import Tracker
+
+
+def run_video(tracker_name, tracker_param, videofile, optional_box=None, debug=None, save_results=False):
+    """Run the tracker on your webcam.
+    args:
+        tracker_name: Name of tracking method.
+        tracker_param: Name of parameter file.
+        debug: Debug level.
+    """
+    tracker = Tracker(tracker_name, tracker_param, "video")
+    tracker.run_video(videofilepath=videofile, optional_box=optional_box, debug=debug, save_results=save_results)
+
+
+def main():
+    parser = argparse.ArgumentParser(description='Run the tracker on your webcam.')
+    parser.add_argument('tracker_name', type=str, help='Name of tracking method.')
+    parser.add_argument('tracker_param', type=str, help='Name of parameter file.')
+    parser.add_argument('videofile', type=str, help='path to a video file.')
+    parser.add_argument('--optional_box', type=float, default=None, nargs="+", help='optional_box with format x y w h.')
+    parser.add_argument('--debug', type=int, default=0, help='Debug level.')
+    parser.add_argument('--save_results', dest='save_results', action='store_true', help='Save bounding boxes')
+    parser.set_defaults(save_results=False)
+
+    args = parser.parse_args()
+
+    run_video(args.tracker_name, args.tracker_param, args.videofile, args.optional_box, args.debug, args.save_results)
+
+
+if __name__ == '__main__':
+    main()