Update app.py

app.py

@@ -14,7 +14,7 @@ with gr.Blocks() as demo:
     * Train/Eval will setup, train, and evaluate the base model
     """)
 
-    def modelTraining():
+    def modelTraining(img):
         mnist = tf.keras.datasets.mnist
 
         (x_train, y_train), (x_test, y_test) = mnist.load_data()
@@ -42,14 +42,13 @@ with gr.Blocks() as demo:
         model.fit(x_train, y_train, epochs=5)
         test_loss, test_acc = model.evaluate(x_test,  y_test, verbose=2)]
 
-        result = "Test accuracy: ", test_acc
+        print "Test accuracy: ", test_acc
 
-        return result
-
-        def predict_image(img):
         # Define any necessary preprocessing steps for the image input here
+        probability_model = tf.keras.Sequential([model, 
+                                         tf.keras.layers.Softmax()])
         # Make a prediction using the model
-        prediction = model.predict(img)
+        prediction = probability_model.predict(img)
 
         # Postprocess the prediction and return it
         return prediction
@@ -57,14 +56,12 @@ with gr.Blocks() as demo:
 
     # Creates the Gradio interface objects
     with gr.Row():
-        with gr.Column(scale=1):
-            submit_btn = gr.Button(value="Train/Eval")
         with gr.Column(scale=2):
+            image_data = gr.Image(label="Upload Image", type="numpy")
+        with gr.Column(scale=1):
             model_performance = gr.Text(label="Model Performance", interactive=False)
             model_prediction = gr.Text(label="Model Prediction", interactive=False)
-            image_data = gr.Image(label="Upload Image", type="numpy")
-    submit_btn.click(modelTraining, [], model_performance)
-    image_data.change(predict_image, image_data, model_prediction)
+    image_data.change(modelTraining, image_data, model_prediction)
     
     
 # creates a local web server