File size: 2,041 Bytes
1bad634 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 |
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, # Image to predict
imgsz=1024, # Prediction image size
conf=0.2, # Confidence threshold
device="cpu" # Device to use (e.g., 'cuda:0' or '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)
|