Update app.py

app.py

@@ -45,13 +45,12 @@ with gr.Blocks() as demo:
         print "Test accuracy: ", test_acc
 
         # Define any necessary preprocessing steps for the image input here
-        probability_model = tf.keras.Sequential([model, 
-                                         tf.keras.layers.Softmax()])
+        probability_model = tf.keras.Sequential([model, tf.keras.layers.Softmax()])
         # Make a prediction using the model
         prediction = probability_model.predict(img)
 
         # Postprocess the prediction and return it
-        return prediction
+        return np.argmax(predictions[0])
         
 
     # Creates the Gradio interface objects