Update app.py

app.py

@@ -204,11 +204,15 @@ import gradio as gr
 import torch
 from PIL import Image
 import torchvision.transforms as T
-import ultralytics
+from ultralytics import YOLO
 
 # Load your model
-model = torch.load("Model_IV.pt")
-model.eval()
+model = YOLO("Model_IV.pt")
+# model = torch.load("Model_IV.pt")
+# model.eval()
+checkpoint = torch.load("model.pt")
+model.load_state_dict(checkpoint)  # Load the saved weights
+model.eval()  # Set the model to evaluation mode
 
 # Define preprocessing
 transform = T.Compose([