JeVoisBase  1.21
JeVois Smart Embedded Machine Vision Toolkit Base Modules
Share this page:
Loading...
Searching...
No Matches
demo.py
Go to the documentation of this file.
1# This file is part of OpenCV Zoo project.
2# It is subject to the license terms in the LICENSE file found in the same directory.
3#
4# Copyright (C) 2021, Shenzhen Institute of Artificial Intelligence and Robotics for Society, all rights reserved.
5# Third party copyrights are property of their respective owners.
6
7import argparse
8
9import numpy as np
10import cv2 as cv
11
12from wechatqrcode import WeChatQRCode
13
14def str2bool(v):
15 if v.lower() in ['on', 'yes', 'true', 'y', 't']:
16 return True
17 elif v.lower() in ['off', 'no', 'false', 'n', 'f']:
18 return False
19 else:
20 raise NotImplementedError
21
22parser = argparse.ArgumentParser(
23 description="WeChat QR code detector for detecting and parsing QR code (https://github.com/opencv/opencv_contrib/tree/master/modules/wechat_qrcode)")
24parser.add_argument('--input', '-i', type=str, help='Path to the input image. Omit for using default camera.')
25parser.add_argument('--detect_prototxt_path', type=str, default='detect_2021sep.prototxt', help='Path to detect.prototxt.')
26parser.add_argument('--detect_model_path', type=str, default='detect_2021sep.caffemodel', help='Path to detect.caffemodel.')
27parser.add_argument('--sr_prototxt_path', type=str, default='sr_2021sep.prototxt', help='Path to sr.prototxt.')
28parser.add_argument('--sr_model_path', type=str, default='sr_2021sep.caffemodel', help='Path to sr.caffemodel.')
29parser.add_argument('--save', '-s', type=str2bool, default=False, help='Set true to save results. This flag is invalid when using camera.')
30parser.add_argument('--vis', '-v', type=str2bool, default=True, help='Set true to open a window for result visualization. This flag is invalid when using camera.')
31args = parser.parse_args()
32
33def visualize(image, res, points, points_color=(0, 255, 0), text_color=(0, 255, 0), fps=None):
34 output = image.copy()
35 h, w, _ = output.shape
36
37 if fps is not None:
38 cv.putText(output, 'FPS: {:.2f}'.format(fps), (0, 15), cv.FONT_HERSHEY_SIMPLEX, 0.5, text_color)
39
40 fontScale = 0.5
41 fontSize = 1
42 for r, p in zip(res, points):
43 p = p.astype(np.int32)
44 for _p in p:
45 cv.circle(output, _p, 10, points_color, -1)
46
47 qrcode_center_x = int((p[0][0] + p[2][0]) / 2)
48 qrcode_center_y = int((p[0][1] + p[2][1]) / 2)
49
50 text_size, baseline = cv.getTextSize(r, cv.FONT_HERSHEY_DUPLEX, fontScale, fontSize)
51 text_x = qrcode_center_x - int(text_size[0] / 2)
52 text_y = qrcode_center_y - int(text_size[1] / 2)
53 cv.putText(output, '{}'.format(r), (text_x, text_y), cv.FONT_HERSHEY_DUPLEX, fontScale, text_color, fontSize)
54
55 return output
56
57
58if __name__ == '__main__':
59 # Instantiate WeChatQRCode
60 model = WeChatQRCode(args.detect_prototxt_path,
61 args.detect_model_path,
62 args.sr_prototxt_path,
63 args.sr_model_path)
64
65 # If input is an image:
66 if args.input is not None:
67 image = cv.imread(args.input)
68 res, points = model.infer(image)
69
70 # Print results:
71 print(res)
72 print(points)
73
74 # Draw results on the input image
75 image = visualize(image, res, points)
76
77 # Save results if save is true
78 if args.save:
79 print('Results saved to result.jpg\n')
80 cv.imwrite('result.jpg', image)
81
82 # Visualize results in a new window
83 if args.vis:
84 cv.namedWindow(args.input, cv.WINDOW_AUTOSIZE)
85 cv.imshow(args.input, image)
86 cv.waitKey(0)
87 else: # Omit input to call default camera
88 deviceId = 0
89 cap = cv.VideoCapture(deviceId)
90
91 tm = cv.TickMeter()
92 while cv.waitKey(1) < 0:
93 hasFrame, frame = cap.read()
94 if not hasFrame:
95 print('No frames grabbed!')
96 break
97
98 # Inference
99 tm.start()
100 res, points = model.infer(frame)
101 tm.stop()
102 fps = tm.getFPS()
103
104 # Draw results on the input image
105 frame = visualize(frame, res, points, fps=fps)
106
107 # Visualize results in a new window
108 cv.imshow('WeChatQRCode Demo', frame)
109
110 tm.reset()
visualize(image, results, box_color=(0, 255, 0), text_color=(0, 0, 255), fps=None)
Definition demo.py:46
int
Definition demo.py:37