libYOLODoc / app.py
offthewallace's picture
Update app.py
e654d05
import gradio as gr
import torch
from PIL import Image
import pdf2image
import numpy as np
model = torch.hub.load('ultralytics/yolov5', 'custom', path='best.pt')
def yolo(im, size=640):
g = (size / max(im.size)) # gain
im = im.resize((int(x * g) for x in im.size), Image.BICUBIC) # resize
results = model(im) # inference
results.render() # updates results.imgs with boxes and labels
return Image.fromarray(results.ims[0])
def detect(pdf):
storelist=[]
path_to_pdf = pdf.name
imgs = pdf2image.convert_from_path(path_to_pdf)
for i in range(len(imgs)):
result = model(np.array(imgs[i]))
result_pred=result.pred
print(result_pred)
for j in range(len(result_pred[0])):
if result_pred[0][j][5]==1 and result_pred[0][j][4]>0.3:
storelist.append(i+1)
if len(storelist)>0:
return "find signature in pdf file page:"+str([*set(storelist)])
else:
return "do not find signature in pdf file"
inputs = gr.inputs.Image(type='pil', label="Original Image")
#outputs = gr.outputs.Image(type="pil", label="Output Image")
title = "Object detection for signature detections in pdf file "
description = "Drop a PDF, it will report pages contain signatures."
gr.Interface(detect, inputs=gr.File(label="PDF"), outputs="label", title=title, description=description,theme="huggingface").launch(debug=True)