未验证 提交 f03283c4 编写于 作者: G George Ni 提交者: GitHub

[cherry-pick][MOT]set mot infer frame rate (#3876)

* add frame_rate for mot infer video

* add doc of frame_rate for mot video infer, test=document_fix
上级 8e80df59
...@@ -262,7 +262,7 @@ Inference a vidoe on single GPU with following command: ...@@ -262,7 +262,7 @@ Inference a vidoe on single GPU with following command:
```bash ```bash
# inference on video and save a video # 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: 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_ ...@@ -273,7 +273,7 @@ CUDA_VISIBLE_DEVICES=0 python tools/infer_mot.py -c configs/mot/fairmot/fairmot_
``` ```
**Notes:** **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 ### 4. Export model
......
...@@ -260,7 +260,7 @@ EvalMOTDataset: ...@@ -260,7 +260,7 @@ EvalMOTDataset:
```bash ```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通过如下命令预测一个图片文件夹,并保存为视频 使用单个GPU通过如下命令预测一个图片文件夹,并保存为视频
...@@ -271,7 +271,7 @@ CUDA_VISIBLE_DEVICES=0 python tools/infer_mot.py -c configs/mot/fairmot/fairmot_ ...@@ -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. 导出预测模型 ### 4. 导出预测模型
......
...@@ -234,6 +234,7 @@ class MOTImageFolder(DetDataset): ...@@ -234,6 +234,7 @@ class MOTImageFolder(DetDataset):
Load MOT dataset with MOT format from image folder or video . Load MOT dataset with MOT format from image folder or video .
Args: Args:
video_file (str): path of the video file, default ''. 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. dataset_dir (str): root directory for dataset.
keep_ori_im (bool): whether to keep original image, default False. keep_ori_im (bool): whether to keep original image, default False.
Set True when used during MOT model inference while saving Set True when used during MOT model inference while saving
...@@ -242,6 +243,7 @@ class MOTImageFolder(DetDataset): ...@@ -242,6 +243,7 @@ class MOTImageFolder(DetDataset):
def __init__(self, def __init__(self,
video_file=None, video_file=None,
frame_rate=-1,
dataset_dir=None, dataset_dir=None,
data_root=None, data_root=None,
image_dir=None, image_dir=None,
...@@ -255,7 +257,7 @@ class MOTImageFolder(DetDataset): ...@@ -255,7 +257,7 @@ class MOTImageFolder(DetDataset):
self.keep_ori_im = keep_ori_im self.keep_ori_im = keep_ori_im
self._imid2path = {} self._imid2path = {}
self.roidbs = None self.roidbs = None
self.frame_rate = 30 self.frame_rate = frame_rate
def check_or_download_dataset(self): def check_or_download_dataset(self):
return return
...@@ -263,17 +265,21 @@ class MOTImageFolder(DetDataset): ...@@ -263,17 +265,21 @@ class MOTImageFolder(DetDataset):
def parse_dataset(self, ): def parse_dataset(self, ):
if not self.roidbs: if not self.roidbs:
if self.video_file is None: if self.video_file is None:
self.frame_rate = 30 # set as default if infer image folder
self.roidbs = self._load_images() self.roidbs = self._load_images()
else: else:
self.roidbs = self._load_video_images() self.roidbs = self._load_video_images()
def _load_video_images(self): def _load_video_images(self):
if self.frame_rate == -1:
# if frame_rate is not set for video, use cv2.VideoCapture
cap = cv2.VideoCapture(self.video_file) cap = cv2.VideoCapture(self.video_file)
self.frame_rate = int(cap.get(cv2.CAP_PROP_FPS)) self.frame_rate = int(cap.get(cv2.CAP_PROP_FPS))
extension = self.video_file.split('.')[-1] extension = self.video_file.split('.')[-1]
output_path = self.video_file.replace('.{}'.format(extension), '') 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( self.video_frames = sorted(
glob.glob(os.path.join(frames_path, '*.png'))) glob.glob(os.path.join(frames_path, '*.png')))
...@@ -334,8 +340,10 @@ class MOTImageFolder(DetDataset): ...@@ -334,8 +340,10 @@ class MOTImageFolder(DetDataset):
self.image_dir = images self.image_dir = images
self.roidbs = self._load_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.video_file = video_file
self.frame_rate = frame_rate
assert os.path.isfile(self.video_file) and _is_valid_video(self.video_file), \ assert os.path.isfile(self.video_file) and _is_valid_video(self.video_file), \
"wrong or unsupported file format: {}".format(self.video_file) "wrong or unsupported file format: {}".format(self.video_file)
self.roidbs = self._load_video_images() self.roidbs = self._load_video_images()
...@@ -345,7 +353,7 @@ def _is_valid_video(f, extensions=('.mp4', '.avi', '.mov', '.rmvb', 'flv')): ...@@ -345,7 +353,7 @@ def _is_valid_video(f, extensions=('.mp4', '.avi', '.mov', '.rmvb', 'flv')):
return f.lower().endswith(extensions) return f.lower().endswith(extensions)
def video2frames(video_path, outpath, **kargs): def video2frames(video_path, outpath, frame_rate, **kargs):
def _dict2str(kargs): def _dict2str(kargs):
cmd_str = '' cmd_str = ''
for k, v in kargs.items(): for k, v in kargs.items():
...@@ -363,7 +371,9 @@ def video2frames(video_path, outpath, **kargs): ...@@ -363,7 +371,9 @@ def video2frames(video_path, outpath, **kargs):
outformat = os.path.join(out_full_path, '%08d.png') outformat = os.path.join(out_full_path, '%08d.png')
cmd = ffmpeg 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) cmd = ''.join(cmd) + _dict2str(kargs)
try: try:
......
...@@ -374,6 +374,7 @@ class Tracker(object): ...@@ -374,6 +374,7 @@ class Tracker(object):
def mot_predict(self, def mot_predict(self,
video_file, video_file,
frame_rate,
image_dir, image_dir,
output_dir, output_dir,
data_type='mot', data_type='mot',
...@@ -401,7 +402,7 @@ class Tracker(object): ...@@ -401,7 +402,7 @@ class Tracker(object):
# run tracking # run tracking
if video_file: if video_file:
seq = video_file.split('/')[-1].split('.')[0] 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)) logger.info('Starting tracking video {}'.format(video_file))
elif image_dir: elif image_dir:
seq = image_dir.split('/')[-1].split('.')[0] seq = image_dir.split('/')[-1].split('.')[0]
...@@ -420,6 +421,7 @@ class Tracker(object): ...@@ -420,6 +421,7 @@ class Tracker(object):
dataloader = create('TestMOTReader')(self.dataset, 0) dataloader = create('TestMOTReader')(self.dataset, 0)
result_filename = os.path.join(result_root, '{}.txt'.format(seq)) result_filename = os.path.join(result_root, '{}.txt'.format(seq))
if frame_rate == -1:
frame_rate = self.dataset.frame_rate frame_rate = self.dataset.frame_rate
with paddle.no_grad(): with paddle.no_grad():
......
...@@ -43,6 +43,11 @@ def parse_args(): ...@@ -43,6 +43,11 @@ def parse_args():
parser = ArgsParser() parser = ArgsParser()
parser.add_argument( parser.add_argument(
'--video_file', type=str, default=None, help='Video name for tracking.') '--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( parser.add_argument(
"--image_dir", "--image_dir",
type=str, type=str,
...@@ -95,6 +100,7 @@ def run(FLAGS, cfg): ...@@ -95,6 +100,7 @@ def run(FLAGS, cfg):
# inference # inference
tracker.mot_predict( tracker.mot_predict(
video_file=FLAGS.video_file, video_file=FLAGS.video_file,
frame_rate=FLAGS.frame_rate,
image_dir=FLAGS.image_dir, image_dir=FLAGS.image_dir,
data_type=cfg.metric.lower(), data_type=cfg.metric.lower(),
model_type=cfg.architecture, model_type=cfg.architecture,
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册