Update app.py

app.py

@@ -11,3 +11,12 @@ 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)
+