ONNX
Zihao Mu commited on
Commit
05dc1e2
·
1 Parent(s): 3d05847

add scale factor to DB demo (#96)

Browse files
Files changed (1) hide show
  1. demo.py +32 -9
demo.py CHANGED
@@ -75,8 +75,12 @@ if __name__ == '__main__':
75
 
76
  # If input is an image
77
  if args.input is not None:
78
- image = cv.imread(args.input)
79
- image = cv.resize(image, [args.width, args.height])
 
 
 
 
80
 
81
  # Inference
82
  results = detector.infer(image)
@@ -86,18 +90,25 @@ if __name__ == '__main__':
86
  recognizer.infer(image, box.reshape(8))
87
  )
88
 
 
 
 
 
 
 
 
89
  # Draw results on the input image
90
- image = visualize(image, results, texts)
91
 
92
  # Save results if save is true
93
  if args.save:
94
  print('Resutls saved to result.jpg\n')
95
- cv.imwrite('result.jpg', image)
96
 
97
  # Visualize results in a new window
98
  if args.vis:
99
  cv.namedWindow(args.input, cv.WINDOW_AUTOSIZE)
100
- cv.imshow(args.input, image)
101
  cv.waitKey(0)
102
  else: # Omit input to call default camera
103
  deviceId = 0
@@ -105,12 +116,17 @@ if __name__ == '__main__':
105
 
106
  tm = cv.TickMeter()
107
  while cv.waitKey(1) < 0:
108
- hasFrame, frame = cap.read()
109
  if not hasFrame:
110
  print('No frames grabbed!')
111
  break
112
 
113
- frame = cv.resize(frame, [args.width, args.height])
 
 
 
 
 
114
  # Inference of text detector
115
  tm.start()
116
  results = detector.infer(frame)
@@ -133,10 +149,17 @@ if __name__ == '__main__':
133
  cv.putText(frame, 'Latency - {}: {:.2f}'.format(recognizer.name, tm.getFPS()), (0, 30), cv.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255))
134
  tm.reset()
135
 
 
 
 
 
 
 
 
136
  # Draw results on the input image
137
- frame = visualize(frame, results, texts)
138
  print(texts)
139
 
140
  # Visualize results in a new Window
141
- cv.imshow('{} Demo'.format(recognizer.name), frame)
142
 
 
75
 
76
  # If input is an image
77
  if args.input is not None:
78
+ original_image = cv.imread(args.input)
79
+ original_w = original_image.shape[1]
80
+ original_h = original_image.shape[0]
81
+ scaleHeight = original_h / args.height
82
+ scaleWidth = original_w / args.width
83
+ image = cv.resize(original_image, [args.width, args.height])
84
 
85
  # Inference
86
  results = detector.infer(image)
 
90
  recognizer.infer(image, box.reshape(8))
91
  )
92
 
93
+ # Scale the results bounding box
94
+ for i in range(len(results[0])):
95
+ for j in range(4):
96
+ box = results[0][i][j]
97
+ results[0][i][j][0] = box[0] * scaleWidth
98
+ results[0][i][j][1] = box[1] * scaleHeight
99
+
100
  # Draw results on the input image
101
+ original_image = visualize(original_image, results, texts)
102
 
103
  # Save results if save is true
104
  if args.save:
105
  print('Resutls saved to result.jpg\n')
106
+ cv.imwrite('result.jpg', original_image)
107
 
108
  # Visualize results in a new window
109
  if args.vis:
110
  cv.namedWindow(args.input, cv.WINDOW_AUTOSIZE)
111
+ cv.imshow(args.input, original_image)
112
  cv.waitKey(0)
113
  else: # Omit input to call default camera
114
  deviceId = 0
 
116
 
117
  tm = cv.TickMeter()
118
  while cv.waitKey(1) < 0:
119
+ hasFrame, original_image = cap.read()
120
  if not hasFrame:
121
  print('No frames grabbed!')
122
  break
123
 
124
+ original_w = original_image.shape[1]
125
+ original_h = original_image.shape[0]
126
+ scaleHeight = original_h / args.height
127
+ scaleWidth = original_w / args.width
128
+
129
+ frame = cv.resize(original_image, [args.width, args.height])
130
  # Inference of text detector
131
  tm.start()
132
  results = detector.infer(frame)
 
149
  cv.putText(frame, 'Latency - {}: {:.2f}'.format(recognizer.name, tm.getFPS()), (0, 30), cv.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255))
150
  tm.reset()
151
 
152
+ # Scale the results bounding box
153
+ for i in range(len(results[0])):
154
+ for j in range(4):
155
+ box = results[0][i][j]
156
+ results[0][i][j][0] = box[0] * scaleWidth
157
+ results[0][i][j][1] = box[1] * scaleHeight
158
+
159
  # Draw results on the input image
160
+ original_image = visualize(original_image, results, texts)
161
  print(texts)
162
 
163
  # Visualize results in a new Window
164
+ cv.imshow('{} Demo'.format(recognizer.name), original_image)
165