Update app.py

app.py

@@ -205,37 +205,58 @@ import torch
 from PIL import Image
 import torchvision.transforms as T
 from ultralytics import YOLO
+import onnxruntime as ort
+import cv2
+import numpy as np
 
 # Load your model
-model = YOLO("Model_IV.pt")
+
+# model = YOLO("Model_IV.pt")
 # model = torch.load("Model_IV.pt")
 # model.eval()
 # checkpoint = torch.load("Model_IV.pt")
 # model.load_state_dict(checkpoint)  # Load the saved weights
 # model.eval()  # Set the model to evaluation mode
 
+# Load the onnx model
+model = ort.InferenceSession("Model_IV.onnx")
+
 # 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
-    results = model(image)
-    # print(type(results))
-    # print(results)
-    annotated_img = results[0].plot()
-    return annotated_img
+    # # Process output (adjust based on your model's format)
+    # results = model(image)
+    # annotated_img = results[0].plot()
+    # return annotated_img
+
+    # Preprocess the image (resize, normalize, etc)
+    input_name = model.get_inputs()[0].name
+    input_shape = model.get_inputs()[0].shape
+
+    # Resize the image to the model's input shape
+    image = cv2.resize(img, (input_shape[2], input_shape[3]))
 
+    # Convert the image to a numpy array and add a batch dimension
+    image = np.expand_dims(img, axis=0)
+
+    # Perform inference
+    output = model.run(None, {input_name: image})
+    print(type(output))
+    print(output)
+
+    return image
+    
 # Gradio interface
 demo = gr.Interface(
     fn=predict,