Update app.py

app.py

@@ -249,6 +249,8 @@ def predict(image):
     # Resize the image to the model's input shape
     image = cv2.resize(image, (input_shape[2], input_shape[3]))
 
+    original_image_shape = image.shape
+
     # Reshape the image to match the model's input shape
     image = image.reshape(3, 640, 640)
 
@@ -280,7 +282,7 @@ def predict(image):
 
     # Reshape the image to match the PIL Image input shape
     print("annotated_img shape before reshape:", annotated_img.shape)
-    annotated_img = annotated_img.reshape(640, 640, 3)
+    annotated_img = annotated_img.reshape(original_image_shape)
     print("annotated_img shape after reshape:", annotated_img.shape)
     
     # annotated_img = (output[0] / 255.0 - mean)/std