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

REPO_ID = "owaiskha9654/Yolov7_Custom_Object_Detection" 
FILENAME = "best.pt"
print(os.getcwd())
yolov7_weights = hf_hub_download(repo_id=REPO_ID, filename=FILENAME)

try:
    model = torch.hub.load('jinfagang/yolov7', 'custom', path=yolov7_weights, force_reload=True)  # local repo
except:
    import zipfile
    l_files = os.listdir("./.cache/torch/hub")
    print(l_files)
    with zipfile.ZipFile("./.cache/torch/hub/main.zip","r") as zip_ref:
        zip_ref.extractall("./.cache/torch/hub/jinfagang_yolov7_main/")
def object_detection(im, size=416):
    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 = "Identificação de Defeitos em Banana"
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=(416, 416), image_mode="RGB", source="upload", label="Image", optional=False)
outputs = gr.outputs.Image(type="pil", label="Output Image")

gr.Interface(
    fn=object_detection,
    inputs=image,
    outputs=outputs,
    title=title,
    description=description,
    examples=[["sample_images/IMG_0125.JPG"], ["sample_images/IMG_0129.JPG"], 
              ["sample_images/IMG_0157.JPG"], ["sample_images/IMG_0158.JPG"]],
).launch()