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 torchkeras import plots 
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/haze_detection.pt')

model = model.to(device)

def load_img (filename):
    img = Image.open(filename).convert("RGB")
    return img

def process_img(image):
    y = image#.to(device)

    with torch.no_grad():
        result = model(y)

        if len(result[0].boxes)>0:
            vis = plots.plot_detection(img,boxes=result[0].boxes.boxes,
                     class_names=class_names, min_score=0.2)
        else:
            vis = img
    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()