Sergey Mikhno
commited on
Commit
·
529ecd0
1
Parent(s):
9efba4e
model
Browse files- app/app.py +15 -0
app/app.py
CHANGED
|
@@ -3,6 +3,9 @@ from doclayout_yolo import YOLOv10
|
|
| 3 |
from huggingface_hub import hf_hub_download
|
| 4 |
import torch
|
| 5 |
import spaces
|
|
|
|
|
|
|
|
|
|
| 6 |
|
| 7 |
import os
|
| 8 |
import json
|
|
@@ -19,6 +22,8 @@ app = Flask(__name__)
|
|
| 19 |
app.secret_key = "super secret key"
|
| 20 |
app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
|
| 21 |
|
|
|
|
|
|
|
| 22 |
|
| 23 |
def allowed_file(filename):
|
| 24 |
return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS
|
|
@@ -54,6 +59,16 @@ def doc_layout():
|
|
| 54 |
blocknames = [names[int(n)] for n in det_res[0].boxes.cls]
|
| 55 |
xyxy = [a.tolist() for a in det_res[0].boxes.xyxy]
|
| 56 |
res = [{"coords": y, "type": x} for x, y in zip(blocknames, xyxy)]
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 57 |
return json.dumps(res)
|
| 58 |
return '''
|
| 59 |
<!doctype html>
|
|
|
|
| 3 |
from huggingface_hub import hf_hub_download
|
| 4 |
import torch
|
| 5 |
import spaces
|
| 6 |
+
import cv2
|
| 7 |
+
from doctr.io import DocumentFile
|
| 8 |
+
from doctr.models import ocr_predictor
|
| 9 |
|
| 10 |
import os
|
| 11 |
import json
|
|
|
|
| 22 |
app.secret_key = "super secret key"
|
| 23 |
app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
|
| 24 |
|
| 25 |
+
doctr_model = ocr_predictor(pretrained=True)
|
| 26 |
+
|
| 27 |
|
| 28 |
def allowed_file(filename):
|
| 29 |
return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS
|
|
|
|
| 59 |
blocknames = [names[int(n)] for n in det_res[0].boxes.cls]
|
| 60 |
xyxy = [a.tolist() for a in det_res[0].boxes.xyxy]
|
| 61 |
res = [{"coords": y, "type": x} for x, y in zip(blocknames, xyxy)]
|
| 62 |
+
pimg = cv2.imread(filename, 0)
|
| 63 |
+
for r in res:
|
| 64 |
+
if r["type"] == "plain text":
|
| 65 |
+
x, y, x1, y1 = r["coords"]
|
| 66 |
+
x, y, x1, y1 = int(x), int(y), int(x1), int(y1)
|
| 67 |
+
pbytes = pimg[y:y1,x:x1].tobytes()
|
| 68 |
+
single_img_doc = DocumentFile.from_images(pbytes)
|
| 69 |
+
result = doctr_model(single_img_doc)
|
| 70 |
+
print(result)
|
| 71 |
+
|
| 72 |
return json.dumps(res)
|
| 73 |
return '''
|
| 74 |
<!doctype html>
|