cleaning up and themeing

.streamlit/config.toml

@@ -0,0 +1,6 @@
+[theme]
+primaryColor = "F36295"
+backgroundColor = "#FFF"
+secondaryBackgroundColor = "#3183D1"
+textColor = "#000"
+font = "sans-serif"

app.py

@@ -20,49 +20,66 @@ from scipy.special import softmax
 @st.cache_resource
 def load_picture():
     """
-    Loads the first 9 images from the mnist dataset and add them to a plot
-    to be displayed in streamlit.
+    Shows the MNIST dataset image
     """
-    # load the mnist dataset
-    (x_train, y_train), (x_test, y_test) = mnist.load_data()
-    # plot the first 9 images
-    for i in range(9):
-        plt.subplot(330 + 1 + i)
-        image = x_train[i] / 255.0
-        plt.imshow(image, cmap=plt.get_cmap("gray"))
-
-    # Save the plot as a png file and show it in streamlit
-    # This is commented out for not because the plot was created and saved in the img directory during the initial run of the app locally
-    # plt.savefig("img/show.png")
     st.image("img/show.png", width=250, caption="First 9 images from the MNIST dataset")
 
 
 def keras_prediction(final, model_path):
+    """Make a predition using a Keras model
+
+    Args:
+        final: The input image
+        model_path: The path of the Keras model to load
+
+    Returns:
+        np.array: Predictions from the model. The probability of each digit.
+        float: Time to make the prediction
+        float: Time to load the model
+    """
+    # load the model
     load_time = time.time()
     model = models.load_model(
         os.path.abspath(os.path.join(os.path.dirname(__file__), model_path))
     )
     after_load_curr = time.time()
+
+    # Make the prediction
     curr_time = time.time()
     prediction = model.predict(final[None, ...])
     after_time = time.time()
+
     return prediction, after_time - curr_time, after_load_curr - load_time
 
 
 def onnx_prediction(final, model_path):
-    im_np = np.expand_dims(final, axis=0)  # Add batch dimension
-    im_np = np.expand_dims(im_np, axis=0)  # Add channel dimension
+    """Make a predition using an Onnx model
+    Args:
+        final: The input image
+        model_path: The path of the Onnx model to load
+
+    Returns:
+        np.array: Predictions from the model. The probability of each digit.
+        float: Time to make the prediction
+        float: Time to load the model
+    """
+    im_np = np.expand_dims(final, axis=0)
+    im_np = np.expand_dims(im_np, axis=0)
     im_np = im_np.astype("float32")
+
+    # Load the model
     load_curr = time.time()
     session = onnxruntime.InferenceSession(model_path, None)
     input_name = session.get_inputs()[0].name
     output_name = session.get_outputs()[0].name
     after_load_curr = time.time()
 
+    # Make the prediction
     curr_time = time.time()
     result = session.run([output_name], {input_name: im_np})
     prediction = softmax(np.array(result).squeeze(), axis=0)
     after_time = time.time()
+
     return prediction, after_time - curr_time, after_load_curr - load_curr
 
 
@@ -70,21 +87,24 @@ def main():
     """
     The main function/primary entry point of the app
     """
-    # write the title of the page as MNIST Digit Recognizer
+    # Setup
+    st.set_page_config(layout="wide")
     st.title("MNIST Digit Recognizer")
 
     col1, col2 = st.columns([0.8, 0.2], gap="small")
     with col1:
         st.markdown(
             """
-            This Streamlit app loads a Keras neural network trained on the MNIST dataset to predict handwritten digits. Draw a digit in the canvas below and see the model's prediction. You can: 
+            This Streamlit app demonstrates the performance of multiple different neural networks (and associated frameworks) trained on the <a href="https://yann.lecun.com/exdb/mnist/">MNIST dataset</a> to predict handwritten digits. Draw a digit in the canvas below and see the model's prediction. You can: 
             - Change the stroke width of the digit using the slider
             - Choose what model you use for predictions
                 - Onnx: The mnist-12 Onnx model from <a href="https://xethub.com/XetHub/onnx-models/src/branch/main/vision/classification/mnist">Onnx's pre-trained MNIST models</a>
                 - Autokeras: A model generated using the <a href="https://autokeras.com/image_classifier/">Autokeras image classifier class</a>
-                - Basic: A simple two layer nueral net where each layer has 300 nodes
+                - Basic: A simple <a href="https://keras.io/">Keras</a> model with two layers where each layer has 300 nodes. The model was trained on the MNIST dataset for 35 epochs.
+            
+            Like any machine learning model, this model is a function of the data it was fed during training. As you can see in the picture, the numbers in the images have a specific shape, location, and size. By playing around with the stroke width and where you draw the digit, you can see how the model's prediction changes.
             
-            Like any machine learning model, this model is a function of the data it was fed during training. As you can see in the picture, the numbers in the images have a specific shape, location, and size. By playing around with the stroke width and where you draw the digit, you can see how the model's prediction changes.""",
+            If you change your selected model after drawing the digit, that same drawing will be used with the newly selected model. To clear your "hand" drawn digit, click the trashcan icon under the drawing canvas.""",
             unsafe_allow_html=True,
         )
     with col2:
@@ -118,8 +138,8 @@ def main():
             background_color="#000",
             background_image=None,
             update_streamlit=True,
-            height=200,
-            width=200,
+            height=300,
+            width=300,
             drawing_mode="freedraw",
             point_display_radius=0,
             key="canvas",
@@ -144,10 +164,14 @@ def main():
                 prediction, pred_time, load_time = onnx_prediction(final, model_path)
 
             # print the prediction
-            st.header(f"Using model: {model_choice}")
-            st.write(f"Prediction: {np.argmax(prediction)}")
-            st.write(f"Load time (in ms): {(load_time) * 1000:.2f}")
-            st.write(f"Prediction time (in ms): {(pred_time) * 1000:.2f}")
+            st.header(f"Results")
+            table_data = {
+                "Model": [model_choice],
+                "Prediction": [np.argmax(prediction)],
+                "Load time (ms)": f"{load_time * 1000:.2f}",
+                "Prediction time (ms)": f"{pred_time * 1000:.2f}",
+            }
+            st.table(table_data)
 
             # Create a 2 column dataframe with one column as the digits and the other as the probability
             data = pd.DataFrame(