sync with remote

app.py

@@ -98,19 +98,35 @@ test_dataset = val.batch(BATCH_SIZE)
 
 st.write(train_dataset)
 
+# function to take a prediction from the model and output an image for display
+def create_mask(pred_mask):
+    pred_mask = tf.argmax(pred_mask, axis=-1)
+    pred_mask = pred_mask[..., tf.newaxis]
+    return pred_mask[0]
+
 def display(display_list):
     fig = plt.figure(figsize=(10, 10))
-    title = ['Input Image', 'Label']
+    title = ['Input Image', 'Label', 'Prediction']  # Updated title list
 
     for i in range(len(display_list)):
         ax = fig.add_subplot(1, len(display_list), i + 1)
         display_resized = tf.reshape(display_list[i], [256, 256])
-        ax.set_title(title[i])
+        ax.set_title(title[i])  # No longer out of range
         ax.imshow(display_resized, cmap='gray')
         ax.axis('off')
 
     st.pyplot(fig)
 
+# helper function to show the image, the label and the prediction
+def show_predictions(dataset=None, num=1):
+    if dataset:
+        for image, label in dataset.take(num):
+            pred_mask = model.predict(image)
+            display([image[0], label[0], create_mask(pred_mask)])
+    else:
+        prediction = create_mask(model.predict(sample_image[tf.newaxis, ...]))
+        display([sample_image, sample_label, prediction])
+
 # Streamlit app interface
 st.title("Cardiac Images Dataset")
 
@@ -143,5 +159,15 @@ model.summary(print_fn=lambda x: model_summary.write(x + '\n'))
 # Display the model summary in Streamlit
 st.markdown(model_summary.getvalue())
 
-# plot the model including the sizes of the model
-tf.keras.utils.plot_model(model, show_shapes=True)
+# Save the model plot
+plot_filename = "model_plot.png"
+tf.keras.utils.plot_model(model, to_file=plot_filename, show_shapes=True)
+
+# Streamlit App
+st.title("Model Architecture")
+
+# Display the model plot
+st.image(plot_filename, caption="Neural Network Architecture", use_container_width=True)
+
+# show a predection, as an example
+show_predictions(test_dataset)