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

[Feature] Support video demo #1014

Merged
merged 6 commits into from
Dec 9, 2021
Merged
Show file tree
Hide file tree
Changes from 3 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
111 changes: 111 additions & 0 deletions demo/video_demo.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,111 @@
from argparse import ArgumentParser

import cv2

from mmseg.apis import inference_segmentor, init_segmentor
from mmseg.core.evaluation import get_palette


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_segmentor(args.config, args.checkpoint, device=args.device)

# build input 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_segmentor(model, frame)

# blend raw image and prediction
draw_img = model.show_result(
frame,
result,
palette=get_palette(args.palette),
show=False,
opacity=args.opacity)

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()
53 changes: 35 additions & 18 deletions docs/get_started.md
Original file line number Diff line number Diff line change
Expand Up @@ -9,24 +9,24 @@

The compatible MMSegmentation and MMCV versions are as below. Please install the correct version of MMCV to avoid installation issues.

| MMSegmentation version | MMCV version |
|:-------------------:|:-------------------:|
| master | mmcv-full>=1.3.13, <1.4.0 |
| 0.19.0 | mmcv-full>=1.3.13, <1.4.0 |
| 0.18.0 | mmcv-full>=1.3.13, <1.4.0 |
| 0.17.0 | mmcv-full>=1.3.7, <1.4.0 |
| 0.16.0 | mmcv-full>=1.3.7, <1.4.0 |
| 0.15.0 | mmcv-full>=1.3.7, <1.4.0 |
| 0.14.1 | mmcv-full>=1.3.7, <1.4.0 |
| 0.14.0 | mmcv-full>=1.3.1, <1.3.2 |
| 0.13.0 | mmcv-full>=1.3.1, <1.3.2 |
| 0.12.0 | mmcv-full>=1.1.4, <1.3.2 |
| 0.11.0 | mmcv-full>=1.1.4, <1.3.0 |
| 0.10.0 | mmcv-full>=1.1.4, <1.3.0 |
| 0.9.0 | mmcv-full>=1.1.4, <1.3.0 |
| 0.8.0 | mmcv-full>=1.1.4, <1.2.0 |
| 0.7.0 | mmcv-full>=1.1.2, <1.2.0 |
| 0.6.0 | mmcv-full>=1.1.2, <1.2.0 |
| MMSegmentation version | MMCV version |
| :--------------------: | :-----------------------: |
| master | mmcv-full>=1.3.13, <1.4.0 |
| 0.19.0 | mmcv-full>=1.3.13, <1.4.0 |
| 0.18.0 | mmcv-full>=1.3.13, <1.4.0 |
| 0.17.0 | mmcv-full>=1.3.7, <1.4.0 |
| 0.16.0 | mmcv-full>=1.3.7, <1.4.0 |
| 0.15.0 | mmcv-full>=1.3.7, <1.4.0 |
| 0.14.1 | mmcv-full>=1.3.7, <1.4.0 |
| 0.14.0 | mmcv-full>=1.3.1, <1.3.2 |
| 0.13.0 | mmcv-full>=1.3.1, <1.3.2 |
| 0.12.0 | mmcv-full>=1.1.4, <1.3.2 |
| 0.11.0 | mmcv-full>=1.1.4, <1.3.0 |
| 0.10.0 | mmcv-full>=1.1.4, <1.3.0 |
| 0.9.0 | mmcv-full>=1.1.4, <1.3.0 |
| 0.8.0 | mmcv-full>=1.1.4, <1.2.0 |
| 0.7.0 | mmcv-full>=1.1.2, <1.2.0 |
| 0.6.0 | mmcv-full>=1.1.2, <1.2.0 |
irvingzhang0512 marked this conversation as resolved.
Show resolved Hide resolved

:::{note}
You need to run `pip uninstall mmcv` first if you have mmcv installed.
Expand Down Expand Up @@ -216,3 +216,20 @@ python demo/image_demo.py demo/demo.jpg configs/pspnet/pspnet_r50-d8_512x1024_40
```

A notebook demo can be found in [demo/inference_demo.ipynb](../demo/inference_demo.ipynb).

Now we also provide a demo script to test a single video.

```shell
irvingzhang0512 marked this conversation as resolved.
Show resolved Hide resolved
python demo/video_demo.py ${VIDEO_FILE} ${CONFIG_FILE} ${CHECKPOINT_FILE} [--device ${DEVICE_NAME}] [--palette-thr ${PALETTE}] \
[--show] [--show-wait-time {SHOW_WAIT_TIME}] [--output-file {OUTPUT_FILE}] [--output-fps {OUTPUT_FPS}] \
[--output-height {OUTPUT_HEIGHT}] [--output-width {OUTPUT_WIDTH}] [--opacity {OPACITY}]
```

Examples:

irvingzhang0512 marked this conversation as resolved.
Show resolved Hide resolved
```shell
# TODO: add mp4 link link `wget -P demo/ `
python demo/video_demo.py demo/demo.mp4 configs/cgnet/cgnet_680x680_60k_cityscapes.py \
checkpoints/cgnet_680x680_60k_cityscapes_20201101_110253-4c0b2f2d.pth \
--device cuda:0 --palette cityscapes --show
```