Spaces:

ssppkenny
/

doclayout

Sleeping

Sergey Mikhno commited on 25 days ago

Commit

ed533ad

1 Parent(s): 8a7e11b

blocks

Files changed (1) hide show

app/app.py CHANGED Viewed

@@ -5,7 +5,7 @@ import torch
 import spaces
 import cv2
 from doctr.io import DocumentFile
-from doctr.models import ocr_predictor
 import os
 import json
@@ -22,7 +22,8 @@ app = Flask(__name__)
 app.secret_key = "super secret key"
 app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
-doctr_model = ocr_predictor(pretrained=True)
 def allowed_file(filename):
@@ -60,17 +61,16 @@ def doc_layout():
             xyxy = [a.tolist() for a in det_res[0].boxes.xyxy]
             res = [{"coords": y, "type": x} for x, y in zip(blocknames, xyxy)]
             pimg = cv2.imread(filename, 0)
-            for r in res:
                 if r["type"] == "plain text":
                     x, y, x1, y1 = r["coords"]
                     x, y, x1, y1 = int(x), int(y), int(x1), int(y1)
-                    block = pimg[y:y1,x:x1]
-                    img_encoded = cv2.imencode('.png', block)[1].tobytes()
-                    single_img_doc = DocumentFile.from_images(img_encoded)
-                    result = doctr_model(single_img_doc)
-                    print(result)
-            return json.dumps(res)
     return '''
     <!doctype html>
     <title>Upload new File</title>

 import spaces
 import cv2
 from doctr.io import DocumentFile
+from doctr.models import detection_predictor
 import os
 import json
 app.secret_key = "super secret key"
 app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
+doctr_model = detection_predictor(arch='db_resnet50', pretrained=True)
 def allowed_file(filename):
             xyxy = [a.tolist() for a in det_res[0].boxes.xyxy]
             res = [{"coords": y, "type": x} for x, y in zip(blocknames, xyxy)]
             pimg = cv2.imread(filename, 0)
+            blocks = dict()
+            for i, r in enumerate(res):
                 if r["type"] == "plain text":
                     x, y, x1, y1 = r["coords"]
                     x, y, x1, y1 = int(x), int(y), int(x1), int(y1)
+                    block = pimg[y:y1, x:x1]
+                    result = doctr_model(block)
+                    blocks["block" + str(i)] = result[0]["words"].tolist()
+            return json.dumps(blocks)
     return '''
     <!doctype html>
     <title>Upload new File</title>