Update app.py

app.py

@@ -281,10 +281,7 @@ def predict(image):
     
     annotated_img = output[0]
 
-    # Reshape the image to match the PIL Image input shape
-    print("annotated_img shape before reshape:", annotated_img.shape)
-    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
     # annotated_img = classes[output[0][0].argmax(0)]
@@ -305,7 +302,12 @@ def predict(image):
     print("Min value of image after normalization:", np.min(annotated_img))
     print("Max value of image after normalization:", np.max(annotated_img))
     print("annotated_img type after normalization:", type(annotated_img))
-    print("annotated_img shape after normalization:", annotated_img.shape)
+    # print("annotated_img shape after normalization:", annotated_img.shape)
+
+    # Reshape the image to match the PIL Image input shape
+    print("annotated_img shape before reshape:", annotated_img.shape)
+    annotated_img = annotated_img.reshape(original_image_shape)
+    print("annotated_img shape after reshape:", annotated_img.shape)
 
     # Convert to PIL Image
     annotated_img = Image.fromarray(annotated_img)