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

yolov7_weights = hf_hub_download(repo_id="LailaMB/visual_pollution_detection", filename="best_640_rpoch56.pt")

model = torch.hub.load('WongKinYiu/yolov7:main', 'custom', yolov7_weights, force_reload=True)  # local repo

def object_detection(im, size=640):
    results = model(im)  # inference
    #results.print() # print results to screen
    #results.show()  # display results
    #results.save()  # save as results1.jpg, results2.jpg... etc.
    results.render() # updates results.imgs with boxes and labels
    return Image.fromarray(results.imgs[0])

title = "visual_pollution_detection"
description = """Esse modelo é uma pequena demonstração baseada em uma análise de cerca de 60 imagens somente. Para resultados mais confiáveis e genéricos, são necessários mais exemplos (imagens).
  """
 
image = gr.inputs.Image(shape=(640, 640), image_mode="RGB", source="upload", label="Imagem", optional=False)
outputs = gr.outputs.Image(type="pil", label="Output Image")

gr.Interface(detect,[gr.Image(type="pil"), gr.Image(type="pil")],description="demo for <a href='https://github.com/WongKinYiu/yolov7' style='text-decoration: underline' target='_blank'>WongKinYiu/yolov7</a> Trainable bag-of-freebies sets new state-of-the-art for real-time object detectors").launch()




#gr.Interface(
#	fn=object_detection,
#    inputs=image,
#    outputs=outputs,
#   title="Visual Pollution Detection",
#    description="demo for <a href='https://github.com/WongKinYiu/yolov7' style='text-decoration: underline' target='_blank'>WongKinYiu/yolov7</a> Trainable bag-of-freebies sets new state-of-the-art for real-time object detectors",
#    examples=[["sample_images/0a1ea4614a9df912eeb8d1b40bffee74.JPG"], ["sample_images/0a2bc0dc2371794509f4b776aff0dd88.JPG"], 
#               ["sample_images/0a4e0e88a05abd96670c8c0c3a67fc73.JPG"], ["sample_images/0a584ddb325ed1ab4083d341280caaa8.JPG"]]
#                ,cache_examples=False).launch()