JeVoisBase  1.18
JeVois Smart Embedded Machine Vision Toolkit Base Modules
Share this page:
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 
7 import sys
8 import argparse
9 
10 import numpy as np
11 import cv2 as cv
12 
13 from crnn import CRNN
14 
15 sys.path.append('../text_detection_db')
16 from db import DB
17 
18 def str2bool(v):
19  if v.lower() in ['on', 'yes', 'true', 'y', 't']:
20  return True
21  elif v.lower() in ['off', 'no', 'false', 'n', 'f']:
22  return False
23  else:
24  raise NotImplementedError
25 
26 backends = [cv.dnn.DNN_BACKEND_OPENCV, cv.dnn.DNN_BACKEND_CUDA]
27 targets = [cv.dnn.DNN_TARGET_CPU, cv.dnn.DNN_TARGET_CUDA, cv.dnn.DNN_TARGET_CUDA_FP16]
28 help_msg_backends = "Choose one of the computation backends: {:d}: OpenCV implementation (default); {:d}: CUDA"
29 help_msg_targets = "Chose one of the target computation devices: {:d}: CPU (default); {:d}: CUDA; {:d}: CUDA fp16"
30 try:
31  backends += [cv.dnn.DNN_BACKEND_TIMVX]
32  targets += [cv.dnn.DNN_TARGET_NPU]
33  help_msg_backends += "; {:d}: TIMVX"
34  help_msg_targets += "; {:d}: NPU"
35 except:
36  print('This version of OpenCV does not support TIM-VX and NPU. Visit https://gist.github.com/fengyuentau/5a7a5ba36328f2b763aea026c43fa45f for more information.')
37 
38 parser = argparse.ArgumentParser(
39  description="An End-to-End Trainable Neural Network for Image-based Sequence Recognition and Its Application to Scene Text Recognition (https://arxiv.org/abs/1507.05717)")
40 parser.add_argument('--input', '-i', type=str, help='Path to the input image. Omit for using default camera.')
41 parser.add_argument('--model', '-m', type=str, default='text_recognition_CRNN_EN_2021sep.onnx', help='Path to the model.')
42 parser.add_argument('--backend', '-b', type=int, default=backends[0], help=help_msg_backends.format(*backends))
43 parser.add_argument('--target', '-t', type=int, default=targets[0], help=help_msg_targets.format(*targets))
44 parser.add_argument('--charset', '-c', type=str, default='charset_36_EN.txt', help='Path to the charset file corresponding to the selected model.')
45 parser.add_argument('--save', '-s', type=str, default=False, help='Set true to save results. This flag is invalid when using camera.')
46 parser.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.')
47 args = parser.parse_args()
48 
49 def visualize(image, boxes, texts, color=(0, 255, 0), isClosed=True, thickness=2):
50  output = image.copy()
51 
52  pts = np.array(boxes[0])
53  output = cv.polylines(output, pts, isClosed, color, thickness)
54  for box, text in zip(boxes[0], texts):
55  cv.putText(output, text, (box[1].astype(np.int32)), cv.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255))
56  return output
57 
58 if __name__ == '__main__':
59  # Instantiate CRNN for text recognition
60  recognizer = CRNN(modelPath=args.model, charsetPath=args.charset)
61  # Instantiate DB for text detection
62  detector = DB(modelPath='../text_detection_db/text_detection_DB_IC15_resnet18_2021sep.onnx',
63  inputSize=[736, 736],
64  binaryThreshold=0.3,
65  polygonThreshold=0.5,
66  maxCandidates=200,
67  unclipRatio=2.0,
68  backendId=args.backend,
69  targetId=args.target
70  )
71 
72  # If input is an image
73  if args.input is not None:
74  image = cv.imread(args.input)
75  image = cv.resize(image, [args.width, args.height])
76 
77  # Inference
78  results = detector.infer(image)
79  texts = []
80  for box, score in zip(results[0], results[1]):
81  texts.append(
82  recognizer.infer(image, box.reshape(8))
83  )
84 
85  # Draw results on the input image
86  image = visualize(image, results, texts)
87 
88  # Save results if save is true
89  if args.save:
90  print('Resutls saved to result.jpg\n')
91  cv.imwrite('result.jpg', image)
92 
93  # Visualize results in a new window
94  if args.vis:
95  cv.namedWindow(args.input, cv.WINDOW_AUTOSIZE)
96  cv.imshow(args.input, image)
97  cv.waitKey(0)
98  else: # Omit input to call default camera
99  deviceId = 0
100  cap = cv.VideoCapture(deviceId)
101 
102  tm = cv.TickMeter()
103  while cv.waitKey(1) < 0:
104  hasFrame, frame = cap.read()
105  if not hasFrame:
106  print('No frames grabbed!')
107  break
108 
109  frame = cv.resize(frame, [736, 736])
110  # Inference of text detector
111  tm.start()
112  results = detector.infer(frame)
113  tm.stop()
114  cv.putText(frame, 'Latency - {}: {:.2f}'.format(detector.name, tm.getFPS()), (0, 15), cv.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255))
115  tm.reset()
116 
117  # Inference of text recognizer
118  if len(results[0]) and len(results[1]):
119  texts = []
120  tm.start()
121  for box, score in zip(results[0], results[1]):
122  result = np.hstack(
123  (box.reshape(8), score)
124  )
125  texts.append(
126  recognizer.infer(frame, box.reshape(8))
127  )
128  tm.stop()
129  cv.putText(frame, 'Latency - {}: {:.2f}'.format(recognizer.name, tm.getFPS()), (0, 30), cv.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255))
130  tm.reset()
131 
132  # Draw results on the input image
133  frame = visualize(frame, results, texts)
134  print(texts)
135 
136  # Visualize results in a new Window
137  cv.imshow('{} Demo'.format(recognizer.name), frame)
138 
demo.str2bool
str2bool
Definition: demo.py:43
demo.visualize
def visualize(image, results, box_color=(0, 255, 0), text_color=(0, 0, 255), fps=None)
Definition: demo.py:46