OpenMMLab Detection Toolbox and Benchmark
https://mmdetection.readthedocs.io/
You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
61 lines
1.9 KiB
61 lines
1.9 KiB
# Copyright (c) OpenMMLab. All rights reserved. |
|
import argparse |
|
|
|
import cv2 |
|
import mmcv |
|
|
|
from mmdet.apis import inference_detector, init_detector |
|
|
|
|
|
def parse_args(): |
|
parser = argparse.ArgumentParser(description='MMDetection 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"') |
|
|
|
model = init_detector(args.config, args.checkpoint, device=args.device) |
|
|
|
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 mmcv.track_iter_progress(video_reader): |
|
result = inference_detector(model, frame) |
|
frame = model.show_result(frame, result, score_thr=args.score_thr) |
|
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()
|
|
|