Update app.py

app.py

@@ -1,4 +1,12 @@
 import streamlit as st
+# THIS MUST BE THE FIRST STREAMLIT COMMAND
+st.set_page_config(
+    page_title="Pet Segmentation with SegFormer",
+    page_icon="🐶",
+    layout="wide",
+    initial_sidebar_state="expanded"
+)
+
 import tensorflow as tf
 from tensorflow.keras import backend
 import numpy as np
@@ -23,13 +31,6 @@ try:
 except Exception as e:
     st.sidebar.error(f"GPU config error: {e}")
 
-st.set_page_config(
-    page_title="Pet Segmentation with SegFormer",
-    page_icon="🐶",
-    layout="wide",
-    initial_sidebar_state="expanded"
-)
-
 # Constants for image preprocessing
 IMAGE_SIZE = 512
 OUTPUT_SIZE = 128
@@ -281,6 +282,54 @@ def create_overlay(image, mask, alpha=0.5):
     
     return overlay
 
+def display_results_side_by_side(original_image, ground_truth_mask=None, predicted_mask=None):
+    """
+    Display results in a side-by-side format similar to colab_code.py
+    
+    Args:
+        original_image: Original input image
+        ground_truth_mask: Optional ground truth segmentation mask
+        predicted_mask: Predicted segmentation mask
+    """
+    # Determine how many images to display
+    cols = 1 + (ground_truth_mask is not None) + (predicted_mask is not None)
+    
+    # Create a figure with multiple columns
+    st.write("### Segmentation Results Comparison")
+    
+    col_list = st.columns(cols)
+    
+    # Display original image
+    with col_list[0]:
+        st.markdown("**Original Image**")
+        st.image(original_image, use_column_width=True)
+    
+    # Display ground truth if available
+    if ground_truth_mask is not None:
+        with col_list[1]:
+            st.markdown("**Ground Truth Mask**")
+            
+            # Colorize ground truth if needed
+            if len(ground_truth_mask.shape) == 2:
+                gt_display = colorize_mask(ground_truth_mask)
+            else:
+                gt_display = ground_truth_mask
+                
+            st.image(gt_display, use_column_width=True)
+    
+    # Display prediction
+    if predicted_mask is not None:
+        with col_list[2 if ground_truth_mask is not None else 1]:
+            st.markdown("**Predicted Mask**")
+            
+            # Colorize prediction if needed
+            if len(predicted_mask.shape) == 2:
+                pred_display = colorize_mask(predicted_mask)
+            else:
+                pred_display = predicted_mask
+                
+            st.image(pred_display, use_column_width=True)
+
 def main():
     st.title("🐶 Pet Segmentation with SegFormer")
     st.markdown("""
@@ -416,6 +465,9 @@ def main():
             # Calculate IoU if ground truth is uploaded
             if uploaded_mask is not None:
                 try:
+                    # Reset the file pointer to the beginning
+                    uploaded_mask.seek(0)
+                    
                     # Read the mask file
                     mask_data = uploaded_mask.read()
                     mask_io = io.BytesIO(mask_data)
@@ -443,10 +495,39 @@ def main():
                     with col3:
                         fg_iou = calculate_iou(gt_mask, resized_mask, 2)
                         st.metric("Foreground IoU", f"{fg_iou:.4f}")
+                        
+                    # For display (original size)
+                    # Reset the file pointer again
+                    uploaded_mask.seek(0)
+                    mask_data = uploaded_mask.read()
+                    mask_io = io.BytesIO(mask_data)
+                    gt_mask_for_display = np.array(Image.open(mask_io))
+                    
+                    # Side-by-side display
+                    display_results_side_by_side(
+                        original_img,
+                        ground_truth_mask=gt_mask_for_display,
+                        predicted_mask=colorized_mask
+                    )
+                    
                 except Exception as e:
                     st.error(f"Error processing ground truth mask: {e}")
                     st.write("Please ensure the mask is valid and has the correct format.")
-            
+                    
+                    # Even with an error, try to display results without the ground truth
+                    display_results_side_by_side(
+                        original_img,
+                        ground_truth_mask=None,
+                        predicted_mask=colorized_mask
+                    )
+            else:
+                # No ground truth, just display original and prediction
+                display_results_side_by_side(
+                    original_img,
+                    ground_truth_mask=None,
+                    predicted_mask=colorized_mask
+                )
+
             # Download buttons
             col1, col2 = st.columns(2)