napatswift commited on
Commit
c3a10c9
·
1 Parent(s): 7831c58
Files changed (1) hide show
  1. main.py +5 -0
main.py CHANGED
@@ -1,6 +1,7 @@
1
  from mmocr.ocr import MMOCR
2
  import gradio as gr
3
  import cv2
 
4
 
5
  ocr = MMOCR(det_config='model/config.py',
6
  det_ckpt='model/epoch_40.pth', device='cpu')
@@ -17,10 +18,13 @@ def get_rec(points):
17
 
18
  def predict(image_input, score_threshold):
19
  draw_img = image_input.copy()
 
 
20
  try:
21
  output = ocr.readtext(image_input)
22
  except:
23
  output = dict(det_polygons=[], det_scores=[])
 
24
  polygons = output['det_polygons']
25
  scores = output['det_scores']
26
  for polygon, score in zip(polygons, scores):
@@ -28,6 +32,7 @@ def predict(image_input, score_threshold):
28
  continue
29
  p0, p1 = get_rec([int(i) for i in polygon])
30
  draw_img = cv2.rectangle(draw_img, p0, p1, (255,0,0), 2)
 
31
  return draw_img
32
 
33
  def run():
 
1
  from mmocr.ocr import MMOCR
2
  import gradio as gr
3
  import cv2
4
+ import sys
5
 
6
  ocr = MMOCR(det_config='model/config.py',
7
  det_ckpt='model/epoch_40.pth', device='cpu')
 
18
 
19
  def predict(image_input, score_threshold):
20
  draw_img = image_input.copy()
21
+ print('image shape', image_input.shape, file=sys.stderr)
22
+
23
  try:
24
  output = ocr.readtext(image_input)
25
  except:
26
  output = dict(det_polygons=[], det_scores=[])
27
+
28
  polygons = output['det_polygons']
29
  scores = output['det_scores']
30
  for polygon, score in zip(polygons, scores):
 
32
  continue
33
  p0, p1 = get_rec([int(i) for i in polygon])
34
  draw_img = cv2.rectangle(draw_img, p0, p1, (255,0,0), 2)
35
+
36
  return draw_img
37
 
38
  def run():