import gradio as gr
import PIL.Image as Image
from ultralytics import YOLO
from huggingface_hub import hf_hub_download
import os

REPO_ID = 'vaivTA/yolov8x_doclaynet'
FILENAME = "weights/best.pt"

print("downloading model...")
model_file = hf_hub_download(repo_id=REPO_ID, filename=FILENAME)
    
model = YOLO(model_file)

def infer(img_path):
    results = model(img_path)
    
    im_array = results[0].plot()
    im = Image.fromarray(im_array[..., ::-1])
    return im
        
    
article = "**이미지를 업로드하세요.**" \
    
demo = gr.Interface(fn=infer,
             inputs=gr.Image(type="numpy", label="Input Image"),
             outputs=gr.Image(type="pil", label="Result Image")
            #  examples=[image_path,]
             )

demo.launch()