Update app.py

app.py

@@ -268,14 +268,12 @@ def predict(image):
     print("type output:", type(output))
     print(output)
 
+    annotated_img = output[0]
     # annotated_img = (output[0] / 255.0 - mean)/std
-    annotated_img = classes[output[0][0].argmax(0)]
-    
+    # annotated_img = classes[output[0][0].argmax(0)]
     print("type annotated image:", type(annotated_img))
     print(annotated_img)
 
-    # annotated_img = output[0]
-
     return annotated_img
     
 # Gradio interface