Update app.py

app.py

@@ -244,17 +244,20 @@ def predict(image):
     input_name = model.get_inputs()[0].name
     input_shape = model.get_inputs()[0].shape
 
-    print("after model.get_inputs()[0].shape:", input_shape)
+    # print("after model.get_inputs()[0].shape:", image.dims)
 
     # Resize the image to the model's input shape
     image = cv2.resize(image, (input_shape[2], input_shape[3]))
     
-    print("after cv2.resize:", input_shape)
+    # print("after cv2.resize:", input_shape)
 
     # Convert the image to a numpy array and add a batch dimension
-    image = np.expand_dims(image, axis=0).astype(np.float32)
+    if len(input_shape) == 4 and input_shape[0] == 1:
+        image = np.expand_dims(image, axis=0)
+    
+    image = image.astype(np.float32)
 
-    print("np.expand_dims:", input_shape)
+    # print("np.expand_dims:", input_shape)
 
     # Perform inference
     output = model.run(None, {input_name: image})