Spaces:
Runtime error
Runtime error
File size: 3,830 Bytes
3b96cb1 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 |
# Copyright (c) OpenMMLab. All rights reserved.
from argparse import ArgumentParser
import cv2
from mmengine.model.utils import revert_sync_batchnorm
from mmseg.apis import inference_model, init_model
from mmseg.apis.inference import show_result_pyplot
def main():
parser = ArgumentParser()
parser.add_argument('video', help='Video file or webcam id')
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(
'--palette',
default='cityscapes',
help='Color palette used for segmentation map')
parser.add_argument(
'--show', action='store_true', help='Whether to show draw result')
parser.add_argument(
'--show-wait-time', default=1, type=int, help='Wait time after imshow')
parser.add_argument(
'--output-file', default=None, type=str, help='Output video file path')
parser.add_argument(
'--output-fourcc',
default='MJPG',
type=str,
help='Fourcc of the output video')
parser.add_argument(
'--output-fps', default=-1, type=int, help='FPS of the output video')
parser.add_argument(
'--output-height',
default=-1,
type=int,
help='Frame height of the output video')
parser.add_argument(
'--output-width',
default=-1,
type=int,
help='Frame width of the output video')
parser.add_argument(
'--opacity',
type=float,
default=0.5,
help='Opacity of painted segmentation map. In (0, 1] range.')
args = parser.parse_args()
assert args.show or args.output_file, \
'At least one output should be enabled.'
# build the model from a config file and a checkpoint file
model = init_model(args.config, args.checkpoint, device=args.device)
if args.device == 'cpu':
model = revert_sync_batchnorm(model)
# build input video
if args.video.isdigit():
args.video = int(args.video)
cap = cv2.VideoCapture(args.video)
assert (cap.isOpened())
input_height = cap.get(cv2.CAP_PROP_FRAME_HEIGHT)
input_width = cap.get(cv2.CAP_PROP_FRAME_WIDTH)
input_fps = cap.get(cv2.CAP_PROP_FPS)
# init output video
writer = None
output_height = None
output_width = None
if args.output_file is not None:
fourcc = cv2.VideoWriter_fourcc(*args.output_fourcc)
output_fps = args.output_fps if args.output_fps > 0 else input_fps
output_height = args.output_height if args.output_height > 0 else int(
input_height)
output_width = args.output_width if args.output_width > 0 else int(
input_width)
writer = cv2.VideoWriter(args.output_file, fourcc, output_fps,
(output_width, output_height), True)
# start looping
try:
while True:
flag, frame = cap.read()
if not flag:
break
# test a single image
result = inference_model(model, frame)
# blend raw image and prediction
draw_img = show_result_pyplot(model, frame, result)
if args.show:
cv2.imshow('video_demo', draw_img)
cv2.waitKey(args.show_wait_time)
if writer:
if draw_img.shape[0] != output_height or draw_img.shape[
1] != output_width:
draw_img = cv2.resize(draw_img,
(output_width, output_height))
writer.write(draw_img)
finally:
if writer:
writer.release()
cap.release()
if __name__ == '__main__':
main()
|