|
|
from flask import Flask |
|
|
from doclayout_yolo import YOLOv10 |
|
|
from huggingface_hub import hf_hub_download |
|
|
|
|
|
import os |
|
|
import json |
|
|
from flask import Flask, flash, request, redirect, url_for |
|
|
from werkzeug.utils import secure_filename |
|
|
|
|
|
UPLOAD_FOLDER = 'upload' |
|
|
ALLOWED_EXTENSIONS = {'txt', 'pdf', 'png', 'jpg', 'jpeg', 'gif'} |
|
|
|
|
|
filepath = hf_hub_download(repo_id="juliozhao/DocLayout-YOLO-DocStructBench", filename="doclayout_yolo_docstructbench_imgsz1024.pt") |
|
|
model = YOLOv10(filepath) |
|
|
|
|
|
app = Flask(__name__) |
|
|
app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER |
|
|
|
|
|
|
|
|
def allowed_file(filename): |
|
|
return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS |
|
|
|
|
|
|
|
|
@app.route("/", methods=['GET', 'POST']) |
|
|
def doc_layout(): |
|
|
if request.method == 'POST': |
|
|
if 'file' not in request.files: |
|
|
flash('no file part') |
|
|
return redirect(request.url) |
|
|
file = request.files['file'] |
|
|
if file.filename == '': |
|
|
flash('No selected file') |
|
|
return redirect(request.url) |
|
|
if file and allowed_file(file.filename): |
|
|
filename = secure_filename(file.filename) |
|
|
file.save(filename) |
|
|
det_res = model.predict( |
|
|
filename, |
|
|
imgsz=1024, |
|
|
conf=0.2, |
|
|
device="cpu" |
|
|
) |
|
|
names = det_res[0].names |
|
|
blocknames = [names[int(n)] for n in det_res[0].boxes.cls] |
|
|
xyxy = [a.tolist() for a in det_res[0].boxes.xyxy] |
|
|
res = [{"coords": y, "type": x} for x, y in zip(blocknames, xyxy)] |
|
|
return json.dumps(res) |
|
|
return ''' |
|
|
<!doctype html> |
|
|
<title>Upload new File</title> |
|
|
<h1>Upload new File</h1> |
|
|
<form method=post enctype=multipart/form-data> |
|
|
<input type=file name=file> |
|
|
<input type=submit value=Upload> |
|
|
</form> |
|
|
''' |
|
|
|
|
|
|
|
|
if __name__ == "__main__": |
|
|
app.run(host='0.0.0.0', port=8080) |
|
|
|
