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
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/yolov8m.pt')

checkpoints = torch.load('Models/haze_detection.pt', map_location=device)

model.load_state_dict(checkpoints['model_state_dict'])

model = model.to(device)

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

def process_img(image):
    img = np.array(image)
    img = img / 255.
    img = img.astype(np.float32)
    y = torch.tensor(img).permute(2,0,1).unsqueeze(0).to(device)

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

    restored_img = result.squeeze().permute(1,2,0).clamp_(0, 1).cpu().detach().numpy()
    restored_img = np.clip(restored_img, 0. , 1.)

    restored_img = (restored_img * 255.0).round().astype(np.uint8)  # float32 to uint8
    return Image.fromarray(restored_img)

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>
'''