Update app.py

app.py

@@ -210,7 +210,7 @@ from ultralytics import YOLO
 model = YOLO("Model_IV.pt")
 # model = torch.load("Model_IV.pt")
 # model.eval()
-# checkpoint = torch.load("Model_IV.pt")
+checkpoint = torch.load("Model_IV.pt")
 # model.load_state_dict(checkpoint)  # Load the saved weights
 # model.eval()  # Set the model to evaluation mode
 
@@ -238,18 +238,18 @@ model = YOLO("Model_IV.pt")
 #     print(checkpoint.keys())
 
 # Define preprocessing
-# transform = T.Compose([
-#     T.Resize((224, 224)),  # Adjust to your model's input size
-#     T.ToTensor(),
-# ])
+transform = T.Compose([
+    T.Resize((224, 224)),  # Adjust to your model's input size
+    T.ToTensor(),
+])
 
 def predict(image):
-    # # Preprocess the image
-    # img_tensor = transform(image).unsqueeze(0)  # Add batch dimension
+    # Preprocess the image
+    img_tensor = transform(image).unsqueeze(0)  # Add batch dimension
     
-    # # Make prediction
-    # with torch.no_grad():
-    #     output = model(img_tensor)
+    # Make prediction
+    with torch.no_grad():
+        output = model(img_tensor)
     
     # Process output (adjust based on your model's format)
     # return output  # or post-process the results as needed