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.
1import argparse
2
3import numpy as np
4import cv2 as cv
5
6from lpd_yunet import LPD_YuNet
7
8def str2bool(v):
9 if v.lower() in ['on', 'yes', 'true', 'y', 't']:
10 return True
11 elif v.lower() in ['off', 'no', 'false', 'n', 'f']:
12 return False
13 else:
14 raise NotImplementedError
15
16backends = [cv.dnn.DNN_BACKEND_OPENCV, cv.dnn.DNN_BACKEND_CUDA]
17targets = [cv.dnn.DNN_TARGET_CPU, cv.dnn.DNN_TARGET_CUDA, cv.dnn.DNN_TARGET_CUDA_FP16]
18help_msg_backends = "Choose one of the computation backends: {:d}: OpenCV implementation (default); {:d}: CUDA"
19help_msg_targets = "Chose one of the target computation devices: {:d}: CPU (default); {:d}: CUDA; {:d}: CUDA fp16"
20try:
21 backends += [cv.dnn.DNN_BACKEND_TIMVX]
22 targets += [cv.dnn.DNN_TARGET_NPU]
23 help_msg_backends += "; {:d}: TIMVX"
24 help_msg_targets += "; {:d}: NPU"
25except:
26 print('This version of OpenCV does not support TIM-VX and NPU. Visit https://gist.github.com/fengyuentau/5a7a5ba36328f2b763aea026c43fa45f for more information.')
27
28parser = argparse.ArgumentParser(description='LPD-YuNet for License Plate Detection')
29parser.add_argument('--input', '-i', type=str, help='Path to the input image. Omit for using default camera.')
30parser.add_argument('--model', '-m', type=str, default='license_plate_detection_lpd_yunet_2022may.onnx', help='Path to the model.')
31parser.add_argument('--backend', '-b', type=int, default=backends[0], help=help_msg_backends.format(*backends))
32parser.add_argument('--target', '-t', type=int, default=targets[0], help=help_msg_targets.format(*targets))
33parser.add_argument('--conf_threshold', type=float, default=0.9, help='Filter out faces of confidence < conf_threshold.')
34parser.add_argument('--nms_threshold', type=float, default=0.3, help='Suppress bounding boxes of iou >= nms_threshold.')
35parser.add_argument('--top_k', type=int, default=5000, help='Keep top_k bounding boxes before NMS.')
36parser.add_argument('--keep_top_k', type=int, default=750, help='Keep keep_top_k bounding boxes after NMS.')
37parser.add_argument('--save', '-s', type=str2bool, default=False, help='Set true to save results. This flag is invalid when using camera.')
38parser.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.')
39args = parser.parse_args()
40
41def visualize(image, dets, line_color=(0, 255, 0), text_color=(0, 0, 255), fps=None):
42 output = image.copy()
43
44 if fps is not None:
45 cv.putText(output, 'FPS: {:.2f}'.format(fps), (0, 15), cv.FONT_HERSHEY_SIMPLEX, 0.5, text_color)
46
47 for det in dets:
48 bbox = det[:-1].astype(np.int32)
49 x1, y1, x2, y2, x3, y3, x4, y4 = bbox
50
51 # Draw the border of license plate
52 cv.line(output, (x1, y1), (x2, y2), line_color, 2)
53 cv.line(output, (x2, y2), (x3, y3), line_color, 2)
54 cv.line(output, (x3, y3), (x4, y4), line_color, 2)
55 cv.line(output, (x4, y4), (x1, y1), line_color, 2)
56
57 return output
58
59if __name__ == '__main__':
60 # Instantiate LPD-YuNet
61 model = LPD_YuNet(modelPath=args.model,
62 confThreshold=args.conf_threshold,
63 nmsThreshold=args.nms_threshold,
64 topK=args.top_k,
65 keepTopK=args.keep_top_k,
66 backendId=args.backend,
67 targetId=args.target)
68
69 # If input is an image
70 if args.input is not None:
71 image = cv.imread(args.input)
72 h, w, _ = image.shape
73
74 # Inference
75 model.setInputSize([w, h])
76 results = model.infer(image)
77
78 # Print results
79 print('{} license plates detected.'.format(results.shape[0]))
80
81 # Draw results on the input image
82 image = visualize(image, results)
83
84 # Save results if save is true
85 if args.save:
86 print('Resutls saved to result.jpg')
87 cv.imwrite('result.jpg', image)
88
89 # Visualize results in a new window
90 if args.vis:
91 cv.namedWindow(args.input, cv.WINDOW_AUTOSIZE)
92 cv.imshow(args.input, image)
93 cv.waitKey(0)
94 else: # Omit input to call default camera
95 deviceId = 0
96 cap = cv.VideoCapture(deviceId)
97 w = int(cap.get(cv.CAP_PROP_FRAME_WIDTH))
98 h = int(cap.get(cv.CAP_PROP_FRAME_HEIGHT))
99 model.setInputSize([w, h])
100
101 tm = cv.TickMeter()
102 while cv.waitKey(1) < 0:
103 hasFrame, frame = cap.read()
104 if not hasFrame:
105 print('No frames grabbed!')
106 break
107
108 # Inference
109 tm.start()
110 results = model.infer(frame) # results is a tuple
111 tm.stop()
112
113 # Draw results on the input image
114 frame = visualize(frame, results, fps=tm.getFPS())
115
116 # Visualize results in a new Window
117 cv.imshow('LPD-YuNet Demo', frame)
118
119 tm.reset()
120