Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

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

Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions configs/mot/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand All @@ -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
Expand Down
4 changes: 2 additions & 2 deletions configs/mot/README_cn.md
Original file line number Diff line number Diff line change
Expand Up @@ -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通过如下命令预测一个图片文件夹,并保存为视频
Expand All @@ -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. 导出预测模型

Expand Down
24 changes: 17 additions & 7 deletions ppdet/data/source/mot.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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,
Expand All @@ -255,25 +257,29 @@ def __init__(self,
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

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')))

Expand Down Expand Up @@ -334,8 +340,10 @@ def set_images(self, images):
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()
Expand All @@ -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():
Expand All @@ -363,7 +371,9 @@ def _dict2str(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:
Expand Down
6 changes: 4 additions & 2 deletions ppdet/engine/tracker.py
Original file line number Diff line number Diff line change
Expand Up @@ -374,6 +374,7 @@ def get_infer_images(self, infer_dir):

def mot_predict(self,
video_file,
frame_rate,
image_dir,
output_dir,
data_type='mot',
Expand Down Expand Up @@ -401,7 +402,7 @@ def mot_predict(self,
# 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]
Expand All @@ -420,7 +421,8 @@ def mot_predict(self,

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']:
Expand Down
6 changes: 6 additions & 0 deletions tools/infer_mot.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down