| |
| import argparse |
|
|
| import cv2 |
| import torch |
|
|
| from mmocr.apis import init_detector, model_inference |
| from mmocr.datasets import build_dataset |
| from mmocr.models import build_detector |
|
|
|
|
| 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=str, default='cuda:0', help='CPU/CUDA device option.') |
| 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() |
|
|
| device = torch.device(args.device) |
|
|
| model = init_detector(args.config, args.checkpoint, device=device) |
|
|
| camera = cv2.VideoCapture(args.camera_id) |
|
|
| print('Press "Esc", "q" or "Q" to exit.') |
| while True: |
| ret_val, img = camera.read() |
| result = model_inference(model, img) |
|
|
| ch = cv2.waitKey(1) |
| if ch == 27 or ch == ord('q') or ch == ord('Q'): |
| break |
|
|
| model.show_result( |
| img, result, score_thr=args.score_thr, wait_time=1, show=True) |
|
|
|
|
| if __name__ == '__main__': |
| main() |
|
|
