app.py

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/dehaze_yolo.pt', map_location=device)

model.load_state_dict(checkpoints['model_state_dict'])

model = model.to(device)