Press key to exit (open-mmlab#1155)

Author: wswday

tools/webcam_demo.py

@@ -0,0 +1,44 @@
+import argparse
+
+import cv2
+import torch
+
+from mmdet.apis import inference_detector, init_detector, show_result
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description='MMDetection webcam demo')
+    parser.add_argument('config', help='test config file path')
+    parser.add_argument('checkpoint', help='checkpoint file')
+    parser.add_argument('--device', type=int, default=0, help='CUDA device id')
+    parser.add_argument(
+        '--camera-id', type=int, default=0, help='camera device id')
+    parser.add_argument(
+        '--score-thr', type=float, default=0.5, help='bbox score threshold')
+    args = parser.parse_args()
+    return args
+
+
+def main():
+    args = parse_args()
+
+    model = init_detector(
+        args.config, args.checkpoint, device=torch.device('cuda', args.device))
+
+    camera = cv2.VideoCapture(args.camera_id)
+
+    print('Press "Esc", "q" or "Q" to exit.')
+    while True:
+        ret_val, img = camera.read()
+        result = inference_detector(model, img)
+
+        ch = cv2.waitKey(1)
+        if ch == 27 or ch == ord('q') or ch == ord('Q'):
+            break
+
+        show_result(
+            img, result, model.CLASSES, score_thr=args.score_thr, wait_time=1)
+
+
+if __name__ == '__main__':
+    main()