Update app.py

app.py

@@ -8,28 +8,27 @@ import numpy as np
 
 # Load the PT model
 model = YOLO("Model_IV.pt")
-checkpoint = torch.load("Model_IV.pt")
+# checkpoint = torch.load("Model_IV.pt")
 
-# Define preprocessing
-transform = T.Compose([
-    T.Resize((224, 224)),  # Adjust to your model's input size
-    T.ToTensor(),
-])
+# # Define preprocessing
+# transform = T.Compose([
+#     T.Resize((224, 224)),  # Adjust to your model's input size
+#     T.ToTensor(),
+# ])
 
 def predict(image):
-    # Preprocess the image by converting the colour space to RGB
+    # Preprocessing: Convert the colour space to RGB
     image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
-    print("converted the colour to RGB.")
+    # print("converted the colour to RGB.")
     
-    # Process output (adjust based on your model's format)
+    # Make prediction 
     results = model(image)
-    print("ran the model")
-    annotated_img = results[0].plot()
-    print("got annotated img")
-    print("type annotated img:", type(annotated_img))
+    #print("ran the model")
 
+    # Postprocessing: Convert the colour space back to BGR
+    annotated_img = results[0].plot()
     annotated_img = cv2.cvtColor(annotated_img, cv2.COLOR_RGB2BGR)
-    print("converted the colour to BGR.")
+    # print("converted the colour to BGR.")
     
     return annotated_img