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)