mirror of https://github.com/open-mmlab/mmyolo.git
[Feature] Add a feature for the video demo (#392)
* Create demo.mp4 * Create video_demo.py * Update large_image_demo.pypull/393/head
parent
446dfbd4b9
commit
ec367dba12
Binary file not shown.
|
@ -2,12 +2,12 @@
|
|||
"""Perform MMYOLO inference on large images (as satellite imagery) as:
|
||||
|
||||
```shell
|
||||
wget -P checkpoint https://download.openmmlab.com/mmyolo/v0/yolov5/yolov5_m-v61_syncbn_fast_8xb16-300e_coco/yolov5_m-v61_syncbn_fast_8xb16-300e_coco_20220917_204944-516a710f.pth syncbn_fast_8xb16-300e_coco/yolov5_m-v61_syncbn_fast_8xb16-300e_coco_20220917_204944-516a710f.pth syncbn_fast_8xb16-300e_coco/yolov5_m-v61_syncbn_fast_8xb16-300e_coco_20220917_204944-516a710f.pth # noqa: E501, E261.
|
||||
wget -P checkpoint https://download.openmmlab.com/mmyolo/v0/yolov5/yolov5_s-v61_syncbn_fast_8xb16-300e_coco/yolov5_s-v61_syncbn_fast_8xb16-300e_coco_20220918_084700-86e02187.pth # noqa: E501, E261.
|
||||
|
||||
python demo/large_image_demo.py \
|
||||
demo/large_image.jpg \
|
||||
configs/yolov5/yolov5_m-v61_syncbn_fast_8xb16-300e_coco.py \
|
||||
checkpoint/yolov5_m-v61_syncbn_fast_8xb16-300e_coco_20220917_204944-516a710f.pth
|
||||
configs/yolov5/yolov5_s-v61_syncbn_fast_8xb16-300e_coco.py \
|
||||
checkpoint/yolov5_s-v61_syncbn_fast_8xb16-300e_coco_20220918_084700-86e02187.pth
|
||||
```
|
||||
"""
|
||||
|
||||
|
|
|
@ -0,0 +1,100 @@
|
|||
# Copyright (c) OpenMMLab. All rights reserved.
|
||||
"""Perform MMYOLO inference on a video as:
|
||||
|
||||
```shell
|
||||
wget -P checkpoint https://download.openmmlab.com/mmyolo/v0/yolov5/yolov5_s-v61_syncbn_fast_8xb16-300e_coco/yolov5_s-v61_syncbn_fast_8xb16-300e_coco_20220918_084700-86e02187.pth # noqa: E501, E261.
|
||||
|
||||
python demo/video_demo.py \
|
||||
demo/video_demo.mp4 \
|
||||
configs/yolov5/yolov5_s-v61_syncbn_fast_8xb16-300e_coco.py \
|
||||
checkpoint/yolov5_s-v61_syncbn_fast_8xb16-300e_coco_20220918_084700-86e02187.pth \
|
||||
--out demo_result.mp4
|
||||
```
|
||||
"""
|
||||
import argparse
|
||||
|
||||
import cv2
|
||||
import mmcv
|
||||
from mmcv.transforms import Compose
|
||||
from mmdet.apis import inference_detector, init_detector
|
||||
from mmengine.utils import track_iter_progress
|
||||
|
||||
from mmyolo.registry import VISUALIZERS
|
||||
from mmyolo.utils import register_all_modules
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(description='MMYOLO video demo')
|
||||
parser.add_argument('video', help='Video file')
|
||||
parser.add_argument('config', help='Config file')
|
||||
parser.add_argument('checkpoint', help='Checkpoint file')
|
||||
parser.add_argument(
|
||||
'--device', default='cuda:0', help='Device used for inference')
|
||||
parser.add_argument(
|
||||
'--score-thr', type=float, default=0.3, help='Bbox score threshold')
|
||||
parser.add_argument('--out', type=str, help='Output video file')
|
||||
parser.add_argument('--show', action='store_true', help='Show video')
|
||||
parser.add_argument(
|
||||
'--wait-time',
|
||||
type=float,
|
||||
default=1,
|
||||
help='The interval of show (s), 0 is block')
|
||||
args = parser.parse_args()
|
||||
return args
|
||||
|
||||
|
||||
def main():
|
||||
args = parse_args()
|
||||
assert args.out or args.show, \
|
||||
('Please specify at least one operation (save/show the '
|
||||
'video) with the argument "--out" or "--show"')
|
||||
|
||||
# register all modules in mmdet into the registries
|
||||
register_all_modules()
|
||||
|
||||
# build the model from a config file and a checkpoint file
|
||||
model = init_detector(args.config, args.checkpoint, device=args.device)
|
||||
|
||||
# build test pipeline
|
||||
model.cfg.test_dataloader.dataset.pipeline[
|
||||
0].type = 'mmdet.LoadImageFromNDArray'
|
||||
test_pipeline = Compose(model.cfg.test_dataloader.dataset.pipeline)
|
||||
|
||||
# init visualizer
|
||||
visualizer = VISUALIZERS.build(model.cfg.visualizer)
|
||||
# the dataset_meta is loaded from the checkpoint and
|
||||
# then pass to the model in init_detector
|
||||
visualizer.dataset_meta = model.dataset_meta
|
||||
|
||||
video_reader = mmcv.VideoReader(args.video)
|
||||
video_writer = None
|
||||
if args.out:
|
||||
fourcc = cv2.VideoWriter_fourcc(*'mp4v')
|
||||
video_writer = cv2.VideoWriter(
|
||||
args.out, fourcc, video_reader.fps,
|
||||
(video_reader.width, video_reader.height))
|
||||
|
||||
for frame in track_iter_progress(video_reader):
|
||||
result = inference_detector(model, frame, test_pipeline=test_pipeline)
|
||||
visualizer.add_datasample(
|
||||
name='video',
|
||||
image=frame,
|
||||
data_sample=result,
|
||||
draw_gt=False,
|
||||
show=False,
|
||||
pred_score_thr=args.score_thr)
|
||||
frame = visualizer.get_image()
|
||||
|
||||
if args.show:
|
||||
cv2.namedWindow('video', 0)
|
||||
mmcv.imshow(frame, 'video', args.wait_time)
|
||||
if args.out:
|
||||
video_writer.write(frame)
|
||||
|
||||
if video_writer:
|
||||
video_writer.release()
|
||||
cv2.destroyAllWindows()
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
Loading…
Reference in New Issue