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

yolov7_custom_weights = "best.pt"

model = torch.hub.load('Owaiskhan9654/yolov7-1:main',model='custom', path_or_model=yolov7_custom_weights, force_reload=True) 

def object_detection(image: gr.inputs.Image = None):
    results = model(image) 
    results.render()  
    count_dict = results.pandas().xyxy[0]['name'].value_counts().to_dict()
    if len(count_dict)>0:
      return Image.fromarray(results.imgs[0]),str(count_dict)
    else:
        return Image.fromarray(results.imgs[0]),'No objects found! Try another image.'

title = "Yolov7 Custom"

inputs = gr.inputs.Image(shape=(1920, 1080), image_mode="RGB", source="upload", label="Upload Image", optional=False)

outputs = gr.outputs.Image(type="pil", label="Output Image")
outputs_cls = gr.Label(label= "Categories Detected Proportion Statistics" )

examples1=[["image0.jpg"],["image1.jpg"],["image2.jpg"],["image3.jpg"],["image4.jpg"],["image5.jpg"],["image6.jpg"],["image7.jpg"]]

Top_Title="Yolov7 🚀 Visual Pollution Detection"
css = ".output-image, .input-image {height: 50rem !important; width: 100% !important;}"
css = ".image-preview {height: auto !important;}"

gr.Interface(
    fn=object_detection,
    inputs=inputs,
    outputs=[outputs,outputs_cls],
    title=Top_Title,
    cache_examples= False,
    allow_flagging='never',
    examples=examples1).launch(debug=True)