Change to cpu

main.py

@@ -5,9 +5,24 @@ import cv2
 model_dir = 'model'
 ocr = MMOCR(det_config=f'{model_dir}/config.py',
             det_ckpt=f'{model_dir}/epoch_40.pth', device='cpu')
-
-def predict(inp):
-    return str(ocr.readtext(inp))
+def get_rec(points):
+    xs = []
+    ys = []
+    for ix, iv in enumerate(points):
+        if ix % 2:
+            ys.append(iv)
+        else:
+            xs.append(iv)
+    return (min(xs), min(ys)), (max(xs), max(ys))
+        
+    
+def predict(image_input):
+    draw_img = image_input.copy()
+    output = ocr.readtext(image_input)
+    for polygon in output['det_poloygons']:
+        p0, p1 = get_rec([int(i) for i in polygon])
+        draw_img = cv2.rectangle(draw_img, p0, p1, (255,255,255))
+    return draw_img
 
 def run():
     demo = gr.Interface(