import gradio as gr 
from PIL import Image
import os
import torch
import torch.nn.functional as F
import torchvision.transforms as transforms
import torchvision
from ultralytics.utils.plotting import Annotator, colors 
import numpy as np
import yaml
from huggingface_hub import hf_hub_download
from ultralytics import YOLO

device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')

model =  YOLO('Models/best.pt')

model = model.to(device)

def load_img (filename):
    if isinstance(img,str):
        img = get_url_img(img) if img.startswith('http') else Image.open(img).convert('RGB')
    return img

def process_img(image):

    with torch.no_grad():
        result = model(source=image)
        lbel=''
        if len(result[0].boxes)>0:
            ann=Annotator(im=image)
            boxes=result[0].boxes
            for element in boxes:
                box=np.array(element.xyxy.cpu()).flatten()
                if element.cls[0].cpu().numpy()==2.0:
                    lbel='car'
                    clr=(0,255,0)
                    
                if element.cls[0].cpu().numpy()==0.0:
                    lbel='bicycle'
                    clr=(255,0,0)
                    
                if element.cls[0].cpu().numpy()==1.0:
                    lbel='bus'
                    clr=(0,0,255)
                    
                if element.cls[0].cpu().numpy()==3.0:
                    lbel='motorcycle'
                    clr=(255,0,255)
                    
                if element.cls[0].cpu().numpy()==4.0:
                    lbel='person'
                    clr=(255,128,0)
                    
                if element.cls[0].cpu().numpy()==5.0:
                    lbel='train'
                    clr=(255,0,128)
                    
                if element.cls[0].cpu().numpy()==6.0:
                    lbel='truck'
                    clr=(0,255,255)
                    
                ann.box_label(box=box, label=lbel, color=clr)
            vis=ann.result()
        else:
            vis = image
    return vis
    
title = "Efficient Hazy Vehicle Detection ✏️🚗🤗"
description = ''' ## [Efficient Hazy Vehicle Detection](https://github.com/cidautai)
[Paula Garrido Mellado](https://github.com/paugar5)
Fundación Cidaut
> **Disclaimer:** please remember this is not a product, thus, you will notice some limitations.
**This demo expects an image with some degradations.**
Due to the GPU memory limitations, the app might crash if you feed a high-resolution image (2K, 4K).
<br>
'''

examples = [['examples/dusttornado.jpg'],
            ['examples/foggy.jpg'], 
            ['examples/haze.jpg'], 
            ["examples/mist.jpg"], 
            ["examples/rain_storm.jpg"],
           ["examples/sand_storm.jpg"],
           ["examples/snow_storm.jpg"]]

css = """
    .image-frame img, .image-container img {
        width: auto;
        height: auto;
        max-width: none;
    }
"""

demo = gr.Interface(
    fn = process_img,
    inputs = [
            gr.Image(type = 'pil', label = 'input')
    ],
    outputs = [gr.Image(type='pil', label = 'output')],
    title = title,
    description = description,
    examples = examples,
    css = css
)

if __name__ == '__main__':
    demo.launch()