import torch
import gradio as gr
from huggingface_hub import hf_hub_download
from PIL import Image

REPO_ID = "michaelwja/burn-detection"
FILENAME = "skin_burn.pt"


yolov7_custom_weights = hf_hub_download(repo_id=REPO_ID, filename=FILENAME,repo_type='space')

model = torch.hub.load('WongKinYiu/yolov7:main',model='custom', path_or_model=yolov7_custom_weights, force_reload=True)
def object_detection(im, size=614):
    results = model(im) 
    results.render()  
    return Image.fromarray(results.imgs[0])

title = "Yolov7 Skin Burn Detection"

image = gr.inputs.Image(shape=(614,614), image_mode="RGB", source="upload", label="Upload Image", optional=False)
outputs = gr.outputs.Image(type="pil", label="Output Image")

Top_Title="Yolov7 Skin Burn Detection | 基于Yolov7的深度学习皮肤烧伤检测模型"
Custom_description="Upload Any Burn Image to Begin. Made by Michael.W"
Footer="北京清华附中计算机高研社团@THIS 2023"


gr.Interface(
    fn=object_detection,
    inputs=image,
    outputs=outputs,
    title=Top_Title,
    description=Custom_description,
    article=Footer,
    examples=[["skin_burn1.jpg"], ["skin_burn2.jpg"]]).launch()