Skip to content

Commit 55fcc1f

Browse files
authored
[MOT] set mot infer frame rate (PaddlePaddle#3875)
* add frame_rate for mot infer video * add doc of frame_rate for mot video infer, test=document_fix
1 parent 2c47839 commit 55fcc1f

File tree

5 files changed

+31
-13
lines changed

5 files changed

+31
-13
lines changed

configs/mot/README.md

+2-2
Original file line numberDiff line numberDiff line change
@@ -262,7 +262,7 @@ Inference a vidoe on single GPU with following command:
262262

263263
```bash
264264
# inference on video and save a video
265-
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
265+
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
266266
```
267267

268268
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_
273273
```
274274

275275
**Notes:**
276-
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`.
276+
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.
277277

278278

279279
### 4. Export model

configs/mot/README_cn.md

+2-2
Original file line numberDiff line numberDiff line change
@@ -260,7 +260,7 @@ EvalMOTDataset:
260260

261261
```bash
262262
# 预测一个视频
263-
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
263+
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
264264
```
265265

266266
使用单个GPU通过如下命令预测一个图片文件夹,并保存为视频
@@ -271,7 +271,7 @@ CUDA_VISIBLE_DEVICES=0 python tools/infer_mot.py -c configs/mot/fairmot/fairmot_
271271
```
272272

273273
**注意:**
274-
请先确保已经安装了[ffmpeg](https://ffmpeg.org/ffmpeg.html), Linux(Ubuntu)平台可以直接用以下命令安装:`apt-get update && apt-get install -y ffmpeg`
274+
请先确保已经安装了[ffmpeg](https://ffmpeg.org/ffmpeg.html), Linux(Ubuntu)平台可以直接用以下命令安装:`apt-get update && apt-get install -y ffmpeg``--frame_rate`表示视频的帧率,表示每秒抽取多少帧,可以自行设置,默认为-1表示会使用OpenCV读取的视频帧率。
275275

276276
### 4. 导出预测模型
277277

ppdet/data/source/mot.py

+17-7
Original file line numberDiff line numberDiff line change
@@ -234,6 +234,7 @@ class MOTImageFolder(DetDataset):
234234
Load MOT dataset with MOT format from image folder or video .
235235
Args:
236236
video_file (str): path of the video file, default ''.
237+
frame_rate (int): frame rate of the video, use cv2 VideoCapture if not set.
237238
dataset_dir (str): root directory for dataset.
238239
keep_ori_im (bool): whether to keep original image, default False.
239240
Set True when used during MOT model inference while saving
@@ -242,6 +243,7 @@ class MOTImageFolder(DetDataset):
242243

243244
def __init__(self,
244245
video_file=None,
246+
frame_rate=-1,
245247
dataset_dir=None,
246248
data_root=None,
247249
image_dir=None,
@@ -255,25 +257,29 @@ def __init__(self,
255257
self.keep_ori_im = keep_ori_im
256258
self._imid2path = {}
257259
self.roidbs = None
258-
self.frame_rate = 30
260+
self.frame_rate = frame_rate
259261

260262
def check_or_download_dataset(self):
261263
return
262264

263265
def parse_dataset(self, ):
264266
if not self.roidbs:
265267
if self.video_file is None:
268+
self.frame_rate = 30 # set as default if infer image folder
266269
self.roidbs = self._load_images()
267270
else:
268271
self.roidbs = self._load_video_images()
269272

270273
def _load_video_images(self):
271-
cap = cv2.VideoCapture(self.video_file)
272-
self.frame_rate = int(cap.get(cv2.CAP_PROP_FPS))
274+
if self.frame_rate == -1:
275+
# if frame_rate is not set for video, use cv2.VideoCapture
276+
cap = cv2.VideoCapture(self.video_file)
277+
self.frame_rate = int(cap.get(cv2.CAP_PROP_FPS))
273278

274279
extension = self.video_file.split('.')[-1]
275280
output_path = self.video_file.replace('.{}'.format(extension), '')
276-
frames_path = video2frames(self.video_file, output_path)
281+
frames_path = video2frames(self.video_file, output_path,
282+
self.frame_rate)
277283
self.video_frames = sorted(
278284
glob.glob(os.path.join(frames_path, '*.png')))
279285

@@ -334,8 +340,10 @@ def set_images(self, images):
334340
self.image_dir = images
335341
self.roidbs = self._load_images()
336342

337-
def set_video(self, video_file):
343+
def set_video(self, video_file, frame_rate):
344+
# update video_file and frame_rate by command line of tools/infer_mot.py
338345
self.video_file = video_file
346+
self.frame_rate = frame_rate
339347
assert os.path.isfile(self.video_file) and _is_valid_video(self.video_file), \
340348
"wrong or unsupported file format: {}".format(self.video_file)
341349
self.roidbs = self._load_video_images()
@@ -345,7 +353,7 @@ def _is_valid_video(f, extensions=('.mp4', '.avi', '.mov', '.rmvb', 'flv')):
345353
return f.lower().endswith(extensions)
346354

347355

348-
def video2frames(video_path, outpath, **kargs):
356+
def video2frames(video_path, outpath, frame_rate, **kargs):
349357
def _dict2str(kargs):
350358
cmd_str = ''
351359
for k, v in kargs.items():
@@ -363,7 +371,9 @@ def _dict2str(kargs):
363371
outformat = os.path.join(out_full_path, '%08d.png')
364372

365373
cmd = ffmpeg
366-
cmd = ffmpeg + [' -i ', video_path, ' -start_number ', ' 0 ', outformat]
374+
cmd = ffmpeg + [
375+
' -i ', video_path, ' -r ', str(frame_rate), ' -f image2 ', outformat
376+
]
367377
cmd = ''.join(cmd) + _dict2str(kargs)
368378

369379
try:

ppdet/engine/tracker.py

+4-2
Original file line numberDiff line numberDiff line change
@@ -374,6 +374,7 @@ def get_infer_images(self, infer_dir):
374374

375375
def mot_predict(self,
376376
video_file,
377+
frame_rate,
377378
image_dir,
378379
output_dir,
379380
data_type='mot',
@@ -401,7 +402,7 @@ def mot_predict(self,
401402
# run tracking
402403
if video_file:
403404
seq = video_file.split('/')[-1].split('.')[0]
404-
self.dataset.set_video(video_file)
405+
self.dataset.set_video(video_file, frame_rate)
405406
logger.info('Starting tracking video {}'.format(video_file))
406407
elif image_dir:
407408
seq = image_dir.split('/')[-1].split('.')[0]
@@ -420,7 +421,8 @@ def mot_predict(self,
420421

421422
dataloader = create('TestMOTReader')(self.dataset, 0)
422423
result_filename = os.path.join(result_root, '{}.txt'.format(seq))
423-
frame_rate = self.dataset.frame_rate
424+
if frame_rate == -1:
425+
frame_rate = self.dataset.frame_rate
424426

425427
with paddle.no_grad():
426428
if model_type in ['JDE', 'FairMOT']:

tools/infer_mot.py

+6
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,11 @@ def parse_args():
4343
parser = ArgsParser()
4444
parser.add_argument(
4545
'--video_file', type=str, default=None, help='Video name for tracking.')
46+
parser.add_argument(
47+
'--frame_rate',
48+
type=int,
49+
default=-1,
50+
help='Video frame rate for tracking.')
4651
parser.add_argument(
4752
"--image_dir",
4853
type=str,
@@ -95,6 +100,7 @@ def run(FLAGS, cfg):
95100
# inference
96101
tracker.mot_predict(
97102
video_file=FLAGS.video_file,
103+
frame_rate=FLAGS.frame_rate,
98104
image_dir=FLAGS.image_dir,
99105
data_type=cfg.metric.lower(),
100106
model_type=cfg.architecture,

0 commit comments

Comments
 (0)