diff --git a/configs/mot/README.md b/configs/mot/README.md index f5debaab4e74f64565f141808ca9bcd403d7dfd8..48823430fa4a2fdc8cc0802af7e4848d923583ab 100644 --- a/configs/mot/README.md +++ b/configs/mot/README.md @@ -262,7 +262,7 @@ Inference a vidoe on single GPU with following command: ```bash # inference on video and save a video -CUDA_VISIBLE_DEVICES=0 python tools/infer_mot.py -c configs/mot/fairmot/fairmot_dla34_30e_1088x608.yml -o weights=https://paddledet.bj.bcebos.com/models/mot/fairmot_dla34_30e_1088x608.pdparams --video_file={your video name}.mp4 --save_videos +CUDA_VISIBLE_DEVICES=0 python tools/infer_mot.py -c configs/mot/fairmot/fairmot_dla34_30e_1088x608.yml -o weights=https://paddledet.bj.bcebos.com/models/mot/fairmot_dla34_30e_1088x608.pdparams --video_file={your video name}.mp4 --frame_rate=20 --save_videos ``` Inference a image folder on single GPU with following command: @@ -273,7 +273,7 @@ CUDA_VISIBLE_DEVICES=0 python tools/infer_mot.py -c configs/mot/fairmot/fairmot_ ``` **Notes:** - Please make sure that [ffmpeg](https://ffmpeg.org/ffmpeg.html) is installed first, on Linux(Ubuntu) platform you can directly install it by the following command:`apt-get update && apt-get install -y ffmpeg`. + Please make sure that [ffmpeg](https://ffmpeg.org/ffmpeg.html) is installed first, on Linux(Ubuntu) platform you can directly install it by the following command:`apt-get update && apt-get install -y ffmpeg`. `--frame_rate` means the frame rate of the video and the frames extracted per second. It can be set by yourself, default value is -1 indicating the video frame rate read by OpenCV. ### 4. Export model diff --git a/configs/mot/README_cn.md b/configs/mot/README_cn.md index ff430d899075be9e398318f821c38e2d09f23bf2..6ff95e3f55033d152b995919b69faacda6ec6966 100644 --- a/configs/mot/README_cn.md +++ b/configs/mot/README_cn.md @@ -260,7 +260,7 @@ EvalMOTDataset: ```bash # 预测一个视频 -CUDA_VISIBLE_DEVICES=0 python tools/infer_mot.py -c configs/mot/fairmot/fairmot_dla34_30e_1088x608.yml -o weights=https://paddledet.bj.bcebos.com/models/mot/fairmot_dla34_30e_1088x608.pdparams --video_file={your video name}.mp4 --save_videos +CUDA_VISIBLE_DEVICES=0 python tools/infer_mot.py -c configs/mot/fairmot/fairmot_dla34_30e_1088x608.yml -o weights=https://paddledet.bj.bcebos.com/models/mot/fairmot_dla34_30e_1088x608.pdparams --video_file={your video name}.mp4 --frame_rate=20 --save_videos ``` 使用单个GPU通过如下命令预测一个图片文件夹，并保存为视频 @@ -271,7 +271,7 @@ CUDA_VISIBLE_DEVICES=0 python tools/infer_mot.py -c configs/mot/fairmot/fairmot_ ``` **注意:** - 请先确保已经安装了[ffmpeg](https://ffmpeg.org/ffmpeg.html), Linux(Ubuntu)平台可以直接用以下命令安装：`apt-get update && apt-get install -y ffmpeg`。 + 请先确保已经安装了[ffmpeg](https://ffmpeg.org/ffmpeg.html), Linux(Ubuntu)平台可以直接用以下命令安装：`apt-get update && apt-get install -y ffmpeg`。`--frame_rate`表示视频的帧率，表示每秒抽取多少帧，可以自行设置，默认为-1表示会使用OpenCV读取的视频帧率。 ### 4. 导出预测模型 diff --git a/ppdet/data/source/mot.py b/ppdet/data/source/mot.py index b5050986f0759333340cd1b526900c0c208f27fa..d163d4ea0ea875fc2f1528e74b6086c6c8543520 100644 --- a/ppdet/data/source/mot.py +++ b/ppdet/data/source/mot.py @@ -234,6 +234,7 @@ class MOTImageFolder(DetDataset): Load MOT dataset with MOT format from image folder or video . Args: video_file (str): path of the video file, default ''. + frame_rate (int): frame rate of the video, use cv2 VideoCapture if not set. dataset_dir (str): root directory for dataset. keep_ori_im (bool): whether to keep original image, default False. Set True when used during MOT model inference while saving @@ -242,6 +243,7 @@ class MOTImageFolder(DetDataset): def __init__(self, video_file=None, + frame_rate=-1, dataset_dir=None, data_root=None, image_dir=None, @@ -255,7 +257,7 @@ class MOTImageFolder(DetDataset): self.keep_ori_im = keep_ori_im self._imid2path = {} self.roidbs = None - self.frame_rate = 30 + self.frame_rate = frame_rate def check_or_download_dataset(self): return @@ -263,17 +265,21 @@ class MOTImageFolder(DetDataset): def parse_dataset(self, ): if not self.roidbs: if self.video_file is None: + self.frame_rate = 30 # set as default if infer image folder self.roidbs = self._load_images() else: self.roidbs = self._load_video_images() def _load_video_images(self): - cap = cv2.VideoCapture(self.video_file) - self.frame_rate = int(cap.get(cv2.CAP_PROP_FPS)) + if self.frame_rate == -1: + # if frame_rate is not set for video, use cv2.VideoCapture + cap = cv2.VideoCapture(self.video_file) + self.frame_rate = int(cap.get(cv2.CAP_PROP_FPS)) extension = self.video_file.split('.')[-1] output_path = self.video_file.replace('.{}'.format(extension), '') - frames_path = video2frames(self.video_file, output_path) + frames_path = video2frames(self.video_file, output_path, + self.frame_rate) self.video_frames = sorted( glob.glob(os.path.join(frames_path, '*.png'))) @@ -334,8 +340,10 @@ class MOTImageFolder(DetDataset): self.image_dir = images self.roidbs = self._load_images() - def set_video(self, video_file): + def set_video(self, video_file, frame_rate): + # update video_file and frame_rate by command line of tools/infer_mot.py self.video_file = video_file + self.frame_rate = frame_rate assert os.path.isfile(self.video_file) and _is_valid_video(self.video_file), \ "wrong or unsupported file format: {}".format(self.video_file) self.roidbs = self._load_video_images() @@ -345,7 +353,7 @@ def _is_valid_video(f, extensions=('.mp4', '.avi', '.mov', '.rmvb', 'flv')): return f.lower().endswith(extensions) -def video2frames(video_path, outpath, **kargs): +def video2frames(video_path, outpath, frame_rate, **kargs): def _dict2str(kargs): cmd_str = '' for k, v in kargs.items(): @@ -363,7 +371,9 @@ def video2frames(video_path, outpath, **kargs): outformat = os.path.join(out_full_path, '%08d.png') cmd = ffmpeg - cmd = ffmpeg + [' -i ', video_path, ' -start_number ', ' 0 ', outformat] + cmd = ffmpeg + [ + ' -i ', video_path, ' -r ', str(frame_rate), ' -f image2 ', outformat + ] cmd = ''.join(cmd) + _dict2str(kargs) try: diff --git a/ppdet/engine/tracker.py b/ppdet/engine/tracker.py index 07fe38800084a0d6e483000c96d7edde60d659f7..15f93427b948826fbad5e73a68dd20ca4df44342 100644 --- a/ppdet/engine/tracker.py +++ b/ppdet/engine/tracker.py @@ -374,6 +374,7 @@ class Tracker(object): def mot_predict(self, video_file, + frame_rate, image_dir, output_dir, data_type='mot', @@ -401,7 +402,7 @@ class Tracker(object): # run tracking if video_file: seq = video_file.split('/')[-1].split('.')[0] - self.dataset.set_video(video_file) + self.dataset.set_video(video_file, frame_rate) logger.info('Starting tracking video {}'.format(video_file)) elif image_dir: seq = image_dir.split('/')[-1].split('.')[0] @@ -420,7 +421,8 @@ class Tracker(object): dataloader = create('TestMOTReader')(self.dataset, 0) result_filename = os.path.join(result_root, '{}.txt'.format(seq)) - frame_rate = self.dataset.frame_rate + if frame_rate == -1: + frame_rate = self.dataset.frame_rate with paddle.no_grad(): if model_type in ['JDE', 'FairMOT']: diff --git a/tools/infer_mot.py b/tools/infer_mot.py index 16d1f193688b12d7fac9b9deaa2336cde95a6462..74b12fd4437044b6f5c254dd8b24e98790c9b39b 100644 --- a/tools/infer_mot.py +++ b/tools/infer_mot.py @@ -43,6 +43,11 @@ def parse_args(): parser = ArgsParser() parser.add_argument( '--video_file', type=str, default=None, help='Video name for tracking.') + parser.add_argument( + '--frame_rate', + type=int, + default=-1, + help='Video frame rate for tracking.') parser.add_argument( "--image_dir", type=str, @@ -95,6 +100,7 @@ def run(FLAGS, cfg): # inference tracker.mot_predict( video_file=FLAGS.video_file, + frame_rate=FLAGS.frame_rate, image_dir=FLAGS.image_dir, data_type=cfg.metric.lower(), model_type=cfg.architecture,