Spaces:
Build error
Build error
napatswift
commited on
Commit
·
c3a10c9
1
Parent(s):
7831c58
Update
Browse files
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():
|