Update app.py

app.py

@@ -75,20 +75,26 @@ def main():
         # Perform inference on the uploaded image
         original_image, mask, segmented_image = perform_inference(image)
 
+        # Display images side by side
+        st.subheader("Image Comparison")
+        fig, axs = plt.subplots(1, 3, figsize=(15, 5))
+
         # Display original image
-        st.subheader("Original Image")
-        fig_original = show_image(original_image, title="Original Image")
-        st.pyplot(fig_original)
+        axs[0].imshow(original_image)
+        axs[0].set_title("Original Image")
+        axs[0].axis('off')
 
         # Display mask
-        st.subheader("Mask")
-        fig_mask = show_image(mask, title="Mask")
-        st.pyplot(fig_mask)
+        axs[1].imshow(mask.squeeze(), cmap='gray')
+        axs[1].set_title("Mask")
+        axs[1].axis('off')
 
         # Display segmented image
-        st.subheader("Segmented Brain Tumor")
-        fig_segmented = show_image(segmented_image, title="Segmented Brain Tumor")
-        st.pyplot(fig_segmented)
+        axs[2].imshow(segmented_image)
+        axs[2].set_title("Segmented Tumor")
+        axs[2].axis('off')
+
+        st.pyplot(fig)
         
 if __name__ == "__main__":
     main()