Update app.py

app.py

@@ -10,6 +10,7 @@ import io
 import gdown
 from transformers import TFSegformerForSemanticSegmentation
 
+# Set page config at the very beginning of the app
 st.set_page_config(
     page_title="Pet Segmentation with SegFormer",
     page_icon="🐶",
@@ -17,25 +18,26 @@ st.set_page_config(
     initial_sidebar_state="expanded"
 )
 
-# Constants for image preprocessing
+# Constants for image preprocessing - matching colab_code.py
 IMAGE_SIZE = 512
 OUTPUT_SIZE = 128
 MEAN = tf.constant([0.485, 0.456, 0.406])
 STD = tf.constant([0.229, 0.224, 0.225])
 
-# Class labels
+# Class labels - DO NOT CHANGE
 ID2LABEL = {0: "background", 1: "border", 2: "foreground/pet"}
 NUM_CLASSES = len(ID2LABEL)
 
 
 @st.cache_resource
 def download_model_from_drive():
+    """Download the model from Google Drive"""
     # Create a models directory
     os.makedirs("models", exist_ok=True)
-    model_path = "models/best_model"
+    model_path = "models/tf_model.h5"
     
     if not os.path.exists(model_path):
-        # Fixed Google Drive URL format for gdown
+        # Correct format for gdown
         url = "https://drive.google.com/file/d/1XObpqG8qZ7YUyiRKbpVvxX11yQSK8Y_3/view?usp=sharing"
         try:
             gdown.download(url, model_path, quiet=False)
@@ -50,12 +52,7 @@ def download_model_from_drive():
 
 @st.cache_resource
 def load_model():
-    """
-    Load the SegFormer model
-    
-    Returns:
-        Loaded model
-    """
+    """Load the SegFormer model"""
     try:
         # First create a base model with the correct architecture
         base_model = TFSegformerForSemanticSegmentation.from_pretrained(
@@ -68,54 +65,35 @@ def load_model():
         
         # Download the trained weights
         model_path = download_model_from_drive()
-        
-        if model_path is not None and os.path.exists(model_path):
-            st.info(f"Loading weights from {model_path}...")
+        if model_path:
             try:
-                # Try to load the weights
                 base_model.load_weights(model_path)
                 st.success("Model weights loaded successfully!")
-                return base_model
             except Exception as e:
                 # st.error(f"Error loading weights: {e}")
-                # st.info("Using base pretrained model instead")
-                return base_model
-        else:
-            st.warning("Using base pretrained model since download failed")
-            return base_model
+                # st.warning("Using base pretrained model instead.")
+        
+        return base_model
             
     except Exception as e:
         st.error(f"Error in load_model: {e}")
-        st.warning("Using default pretrained model")
-        # Fall back to pretrained model as a last resort
-        return TFSegformerForSemanticSegmentation.from_pretrained(
-            "nvidia/mit-b0",
-            num_labels=NUM_CLASSES,
-            id2label=ID2LABEL,
-            label2id={label: id for id, label in ID2LABEL.items()},
-            ignore_mismatched_sizes=True
-        )
+        return None
+
 
 def normalize_image(input_image):
-    """
-    Normalize the input image
-    
-    Args:
-        input_image: Image to normalize
-        
-    Returns:
-        Normalized image
-    """
+    """Normalize image with ImageNet stats"""
     input_image = tf.image.convert_image_dtype(input_image, tf.float32)
     input_image = (input_image - MEAN) / tf.maximum(STD, backend.epsilon())
     return input_image
 
-def preprocess_image(image):
+
+def preprocess_image(image, is_dataset_image=False):
     """
-    Preprocess image for model input
+    Preprocess image exactly like in colab_code.py
     
     Args:
         image: PIL Image to preprocess
+        is_dataset_image: Whether the image is from the Oxford-IIIT Pet dataset
         
     Returns:
         Preprocessed image tensor, original image
@@ -126,8 +104,14 @@ def preprocess_image(image):
     # Store original image for display
     original_img = img_array.copy()
     
-    # Resize to target size
-    img_resized = tf.image.resize(img_array, (IMAGE_SIZE, IMAGE_SIZE))
+    # Resize to target size with preserve_aspect_ratio=False
+    img_resized = tf.image.resize(
+        img_array, 
+        (IMAGE_SIZE, IMAGE_SIZE),
+        method='bilinear',
+        preserve_aspect_ratio=False,  # Ensure exact dimensions
+        antialias=True
+    )
     
     # Normalize
     img_normalized = normalize_image(img_resized)
@@ -140,6 +124,41 @@ def preprocess_image(image):
     
     return img_batch, original_img
 
+
+def process_uploaded_mask(mask_array, from_dataset=True):
+    """
+    Process an uploaded mask from the dataset to match app's format
+    
+    Args:
+        mask_array: Numpy array of the mask
+        from_dataset: Whether the mask is from the original dataset
+        
+    Returns:
+        Processed mask with values 0,1,2
+    """
+    # Handle RGBA images
+    if len(mask_array.shape) == 3 and mask_array.shape[2] == 4:
+        mask_array = mask_array[:,:,:3]
+    
+    # Convert RGB to grayscale if needed
+    if len(mask_array.shape) == 3 and mask_array.shape[2] >= 3:
+        mask_array = cv2.cvtColor(mask_array, cv2.COLOR_RGB2GRAY)
+    
+    # For dataset masks, convert from original values (1,2,3) to app values (0,1,2)
+    if from_dataset:
+        processed_mask = np.zeros_like(mask_array)
+        
+        # Map dataset values to app values
+        processed_mask[mask_array == 1] = 2  # Foreground/pet (1→2)
+        processed_mask[mask_array == 2] = 1  # Border (2→1)
+        processed_mask[mask_array == 3] = 0  # Background (3→0)
+        
+        return processed_mask
+    else:
+        # For non-dataset masks, assume they're already in the right format
+        return mask_array
+
+
 def create_mask(pred_mask):
     """
     Convert model prediction to displayable mask
@@ -150,64 +169,96 @@ def create_mask(pred_mask):
     Returns:
         Processed mask (2D array)
     """
+    # Take argmax along the class dimension
     pred_mask = tf.math.argmax(pred_mask, axis=1)
+    
+    # Remove batch dimension and convert to numpy
     pred_mask = tf.squeeze(pred_mask)
+    
     return pred_mask.numpy()
 
+
 def colorize_mask(mask):
     """
-    Apply colors to segmentation mask
+    Colorize a segmentation mask for visualization
     
     Args:
-        mask: Segmentation mask (2D array)
+        mask: Segmentation mask (2D array with class indices)
         
     Returns:
-        Colorized mask (3D RGB array)
+        Colorized mask (3D array with RGB colors)
     """
-    # Ensure the mask is 2D
-    if len(mask.shape) > 2:
-        mask = np.squeeze(mask)
-    
-    # Define colors for each class (RGB)
+    # Define colors for visualization
     colors = [
-        [0, 0, 0],      # Background (black)
-        [255, 0, 0],    # Border (red)
-        [0, 0, 255]     # Foreground/pet (blue)
+        [0, 0, 0],       # Black for background (0)
+        [255, 255, 0],   # Yellow for border (1)
+        [255, 0, 0]      # Red for foreground/pet (2)
     ]
     
     # Create RGB mask
-    rgb_mask = np.zeros((mask.shape[0], mask.shape[1], 3), dtype=np.uint8)
+    height, width = mask.shape
+    colorized = np.zeros((height, width, 3), dtype=np.uint8)
     
+    # Apply colors
     for i, color in enumerate(colors):
-        class_mask = (mask == i).astype(np.uint8)
-        for c in range(3):
-            rgb_mask[:, :, c] += class_mask * color[c]
+        colorized[mask == i] = color
     
-    return rgb_mask
+    return colorized
+
+
+def create_overlay(image, mask, alpha=0.5):
+    """
+    Create an overlay of mask on original image
+    
+    Args:
+        image: Original image
+        mask: Colorized segmentation mask
+        alpha: Transparency level (0-1)
+        
+    Returns:
+        Overlay image
+    """
+    # Ensure mask shape matches image
+    if image.shape[:2] != mask.shape[:2]:
+        mask = cv2.resize(mask, (image.shape[1], image.shape[0]))
+    
+    # Create blend
+    overlay = cv2.addWeighted(
+        image, 
+        1, 
+        mask.astype(np.uint8), 
+        alpha, 
+        0
+    )
+    
+    return overlay
+
 
 def calculate_iou(y_true, y_pred, class_idx=None):
     """
-    Calculate IoU (Intersection over Union) for segmentation masks
+    Calculate IoU (Intersection over Union)
     
     Args:
-        y_true: Ground truth segmentation mask
-        y_pred: Predicted segmentation mask
-        class_idx: Index of the class to calculate IoU for (None for mean IoU)
+        y_true: Ground truth mask
+        y_pred: Predicted mask
+        class_idx: Class index to compute IoU for (if None, compute mean IoU)
         
     Returns:
         IoU score
     """
     if class_idx is not None:
-        # Binary IoU for specific class
+        # Convert to binary masks for specific class
         y_true_class = (y_true == class_idx).astype(np.float32)
         y_pred_class = (y_pred == class_idx).astype(np.float32)
         
+        # Calculate intersection and union
         intersection = np.sum(y_true_class * y_pred_class)
         union = np.sum(y_true_class) + np.sum(y_pred_class) - intersection
         
-        iou = intersection / (union + 1e-6)
+        # Return IoU score
+        return float(intersection) / float(union) if union > 0 else 0.0
     else:
-        # Mean IoU across all classes
+        # Calculate mean IoU across all classes
         class_ious = []
         for idx in range(NUM_CLASSES):
             class_iou = calculate_iou(y_true, y_pred, idx)
@@ -217,32 +268,105 @@ def calculate_iou(y_true, y_pred, class_idx=None):
     
     return iou
 
-def create_overlay(image, mask, alpha=0.5):
+
+def calculate_dice(y_true, y_pred, class_idx=None):
     """
-    Create an overlay of mask on original image
+    Calculate Dice coefficient (F1 score)
     
     Args:
-        image: Original image
-        mask: Segmentation mask
-        alpha: Transparency level (0-1)
+        y_true: Ground truth mask
+        y_pred: Predicted mask
+        class_idx: Class index to compute Dice for (if None, compute mean Dice)
         
     Returns:
-        Overlay image
+        Dice score
     """
-    # Ensure mask shape matches image
-    if image.shape[:2] != mask.shape[:2]:
-        mask = cv2.resize(mask, (image.shape[1], image.shape[0]))
+    if class_idx is not None:
+        # Convert to binary masks for specific class
+        y_true_class = (y_true == class_idx).astype(np.float32)
+        y_pred_class = (y_pred == class_idx).astype(np.float32)
+        
+        # Calculate intersection and sum of areas
+        intersection = 2.0 * np.sum(y_true_class * y_pred_class)
+        sum_areas = np.sum(y_true_class) + np.sum(y_pred_class)
+        
+        # Return Dice score
+        return float(intersection) / float(sum_areas) if sum_areas > 0 else 0.0
+    else:
+        # Calculate mean Dice across all classes
+        class_dices = []
+        for idx in range(NUM_CLASSES):
+            class_dice = calculate_dice(y_true, y_pred, idx)
+            class_dices.append(class_dice)
+        
+        dice = np.mean(class_dices)
     
-    # Create blend
-    overlay = cv2.addWeighted(
-        image, 
-        1, 
-        mask.astype(np.uint8), 
-        alpha, 
-        0
-    )
+    return dice
+
+
+def calculate_pixel_accuracy(y_true, y_pred):
+    """
+    Calculate pixel accuracy
     
-    return overlay
+    Args:
+        y_true: Ground truth mask
+        y_pred: Predicted mask
+        
+    Returns:
+        Pixel accuracy
+    """
+    correct = np.sum(y_true == y_pred)
+    total = y_true.size
+    return float(correct) / float(total)
+
+
+def display_side_by_side(original_img, gt_mask=None, pred_mask=None, overlay=None):
+    """
+    Display images side by side
+    
+    Args:
+        original_img: Original input image
+        gt_mask: Ground truth segmentation mask (optional)
+        pred_mask: Predicted segmentation mask
+        overlay: Overlay of mask on original image
+    """
+    # Determine number of columns based on available images
+    columns = 1  # Start with original image
+    if gt_mask is not None:
+        columns += 1
+    if pred_mask is not None:
+        columns += 1
+    if overlay is not None:
+        columns += 1
+    
+    cols = st.columns(columns)
+    
+    # Display original image
+    with cols[0]:
+        st.markdown("### Original Image")
+        st.image(original_img, use_column_width=True)
+    
+    # Display ground truth mask if available
+    col_idx = 1
+    if gt_mask is not None:
+        with cols[col_idx]:
+            st.markdown("### Ground Truth Mask")
+            st.image(gt_mask, use_column_width=True)
+        col_idx += 1
+    
+    # Display predicted mask if available
+    if pred_mask is not None:
+        with cols[col_idx]:
+            st.markdown("### Predicted Mask")
+            st.image(pred_mask, use_column_width=True)
+        col_idx += 1
+    
+    # Display overlay if available
+    if overlay is not None:
+        with cols[col_idx]:
+            st.markdown("### Overlay")
+            st.image(overlay, use_column_width=True)
+
 
 def main():
     st.title("🐶 Pet Segmentation with SegFormer")
@@ -254,23 +378,16 @@ def main():
         - **Foreground**: The pet itself
     """)
     
-    # Sidebar
-    st.sidebar.header("Model Information")
-    st.sidebar.markdown("""
-        **SegFormer** is a state-of-the-art semantic segmentation model based on transformers.
-        
-        Key features:
-        - Hierarchical transformer encoder
-        - Lightweight MLP decoder
-        - Efficient mix of local and global attention
-        
-        This implementation uses the MIT-B0 variant fine-tuned on the Oxford-IIIT Pet dataset.
-    """)
+    # Sidebar settings
+    st.sidebar.title("Settings")
+    
+    # Debug mode toggle
+    debug_mode = st.sidebar.checkbox("Debug Mode", value=False)
     
-    # Advanced settings in sidebar
-    st.sidebar.header("Settings")
+    # Dataset image toggle - important for processing Oxford-IIIT Pet masks
+    dataset_image = st.sidebar.checkbox("Image is from Oxford-IIIT Pet dataset", value=True)
     
-    # Overlay opacity
+    # Overlay opacity control
     overlay_opacity = st.sidebar.slider(
         "Overlay Opacity", 
         min_value=0.1, 
@@ -284,7 +401,8 @@ def main():
         model = load_model()
         
     if model is None:
-        st.error("Failed to load model. Using default pretrained model instead.")
+        st.error("Failed to load model. Please check your model path and try again.")
+        return
     else:
         st.sidebar.success("Model loaded successfully!")
     
@@ -300,22 +418,16 @@ def main():
             image_bytes = uploaded_image.read()
             image = Image.open(io.BytesIO(image_bytes))
                 
-            col1, col2 = st.columns(2)
-            
-            with col1:
-                st.subheader("Original Image")
-                st.image(image, caption="Uploaded Image", use_column_width=True)
-            
             # Preprocess and predict
             with st.spinner("Generating segmentation mask..."):
                 # Preprocess the image
-                img_tensor, original_img = preprocess_image(image)
+                img_tensor, original_img = preprocess_image(image, is_dataset_image=dataset_image)
                 
                 # Make prediction
                 outputs = model(pixel_values=img_tensor, training=False)
                 logits = outputs.logits
                 
-                # Create visualization mask
+                # Create mask
                 mask = create_mask(logits)
                 
                 # Colorize the mask
@@ -324,10 +436,86 @@ def main():
                 # Create overlay
                 overlay = create_overlay(original_img, colorized_mask, alpha=overlay_opacity)
             
+            # Prepare for metrics calculation (if ground truth is provided)
+            gt_mask = None
+            gt_mask_colorized = None
+            metrics_calculated = False
+            
+            # Calculate metrics 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)
+                    gt_mask_raw = np.array(Image.open(mask_io))
+                    
+                    if debug_mode:
+                        st.write(f"Ground truth mask shape: {gt_mask_raw.shape}")
+                        st.write(f"Ground truth mask unique values: {np.unique(gt_mask_raw)}")
+                    
+                    # Process the mask based on source
+                    gt_mask = process_uploaded_mask(gt_mask_raw, from_dataset=dataset_image)
+                    
+                    # Colorize for display
+                    gt_mask_colorized = colorize_mask(gt_mask)
+                    
+                    # Resize for comparison
+                    gt_mask_resized = cv2.resize(gt_mask, (mask.shape[0], mask.shape[1]), 
+                                                interpolation=cv2.INTER_NEAREST)
+                    
+                    if debug_mode:
+                        st.write(f"Processed GT mask shape: {gt_mask_resized.shape}")
+                        st.write(f"Processed GT unique values: {np.unique(gt_mask_resized)}")
+                        st.write(f"Prediction mask unique values: {np.unique(mask)}")
+                    
+                    # Calculate metrics
+                    iou_score = calculate_iou(gt_mask_resized, mask)
+                    dice_score = calculate_dice(gt_mask_resized, mask)
+                    accuracy = calculate_pixel_accuracy(gt_mask_resized, mask)
+                    
+                    metrics_calculated = True
+                except Exception as e:
+                    st.error(f"Error processing ground truth mask: {e}")
+                    if debug_mode:
+                        import traceback
+                        st.code(traceback.format_exc())
+            
             # Display results
-            with col2:
-                st.subheader("Segmentation Result")
-                st.image(overlay, caption="Segmentation Overlay", use_column_width=True)
+            display_side_by_side(
+                original_img,
+                gt_mask_colorized,
+                colorized_mask,
+                overlay
+            )
+            
+            # Display metrics if calculated
+            if metrics_calculated:
+                st.header("Segmentation Metrics")
+                
+                # Display overall metrics
+                col1, col2, col3 = st.columns(3)
+                with col1:
+                    st.metric("Mean IoU", f"{iou_score:.4f}")
+                with col2:
+                    st.metric("Mean Dice", f"{dice_score:.4f}")
+                with col3:
+                    st.metric("Pixel Accuracy", f"{accuracy:.4f}")
+                
+                # Display class-specific metrics
+                st.subheader("Metrics by Class")
+                cols = st.columns(NUM_CLASSES)
+                class_names = ["Background", "Border", "Foreground/Pet"]
+                
+                for i, (col, name) in enumerate(zip(cols, class_names)):
+                    with col:
+                        st.markdown(f"**{name}**")
+                        class_iou = calculate_iou(gt_mask_resized, mask, i)
+                        class_dice = calculate_dice(gt_mask_resized, mask, i)
+                        st.metric("IoU", f"{class_iou:.4f}")
+                        st.metric("Dice", f"{class_dice:.4f}")
             
             # Display segmentation details
             st.header("Segmentation Details")
@@ -351,86 +539,68 @@ def main():
                 mask_fg = np.where(mask == 2, 255, 0).astype(np.uint8)
                 st.image(mask_fg, caption="Foreground", use_column_width=True)
             
-            # Calculate IoU if ground truth is uploaded
-            if uploaded_mask is not None:
-                try:
-                    # Read the mask file
-                    mask_data = uploaded_mask.read()
-                    mask_io = io.BytesIO(mask_data)
-                    gt_mask = np.array(Image.open(mask_io).resize((OUTPUT_SIZE, OUTPUT_SIZE), Image.NEAREST))
-                    
-                    # Handle different mask formats
-                    if len(gt_mask.shape) == 3 and gt_mask.shape[2] == 3:
-                        # Convert RGB to single channel if needed
-                        gt_mask = cv2.cvtColor(gt_mask, cv2.COLOR_RGB2GRAY)
-                    
-                    # Calculate and display IoU
-                    resized_mask = cv2.resize(mask, (OUTPUT_SIZE, OUTPUT_SIZE), interpolation=cv2.INTER_NEAREST)
-                    iou_score = calculate_iou(gt_mask, resized_mask)
-                    st.success(f"Mean IoU: {iou_score:.4f}")
-                    
-                    # Display specific class IoUs
-                    st.markdown("### IoU by Class")
-                    col1, col2, col3 = st.columns(3)
-                    with col1:
-                        bg_iou = calculate_iou(gt_mask, resized_mask, 0)
-                        st.metric("Background IoU", f"{bg_iou:.4f}")
-                    with col2:
-                        border_iou = calculate_iou(gt_mask, resized_mask, 1)
-                        st.metric("Border IoU", f"{border_iou:.4f}")
-                    with col3:
-                        fg_iou = calculate_iou(gt_mask, resized_mask, 2)
-                        st.metric("Foreground IoU", f"{fg_iou:.4f}")
-                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.")
-            
             # Download buttons
-            col1, col2 = st.columns(2)
+            st.header("Download Results")
+            col1, col2, col3 = st.columns(3)
             
             with col1:
-                # Convert mask to PNG for download
-                mask_colored = Image.fromarray(colorized_mask)
-                mask_bytes = io.BytesIO()
-                mask_colored.save(mask_bytes, format='PNG')
-                mask_bytes = mask_bytes.getvalue()
+                # Download prediction as PNG
+                pred_pil = Image.fromarray(colorized_mask)
+                pred_bytes = io.BytesIO()
+                pred_pil.save(pred_bytes, format='PNG')
+                pred_bytes = pred_bytes.getvalue()
                 
                 st.download_button(
-                    label="Download Segmentation Mask",
-                    data=mask_bytes,
-                    file_name="pet_segmentation_mask.png",
+                    label="Download Prediction",
+                    data=pred_bytes,
+                    file_name="prediction.png",
                     mime="image/png"
                 )
             
             with col2:
-                # Convert overlay to PNG for download
-                overlay_img = Image.fromarray(overlay)
+                # Download overlay as PNG
+                overlay_pil = Image.fromarray(overlay)
                 overlay_bytes = io.BytesIO()
-                overlay_img.save(overlay_bytes, format='PNG')
+                overlay_pil.save(overlay_bytes, format='PNG')
                 overlay_bytes = overlay_bytes.getvalue()
                 
                 st.download_button(
-                    label="Download Overlay Image",
+                    label="Download Overlay",
                     data=overlay_bytes,
-                    file_name="pet_segmentation_overlay.png",
+                    file_name="overlay.png",
                     mime="image/png"
                 )
+            
+            if metrics_calculated:
+                with col3:
+                    # Create CSV with metrics
+                    metrics_csv = f"Metric,Overall,Background,Border,Foreground\n"
+                    metrics_csv += f"IoU,{iou_score:.4f},{calculate_iou(gt_mask_resized, mask, 0):.4f},{calculate_iou(gt_mask_resized, mask, 1):.4f},{calculate_iou(gt_mask_resized, mask, 2):.4f}\n"
+                    metrics_csv += f"Dice,{dice_score:.4f},{calculate_dice(gt_mask_resized, mask, 0):.4f},{calculate_dice(gt_mask_resized, mask, 1):.4f},{calculate_dice(gt_mask_resized, mask, 2)::.4f}\n"
+                    metrics_csv += f"Accuracy,{accuracy:.4f},,,"
+                    
+                    st.download_button(
+                        label="Download Metrics",
+                        data=metrics_csv,
+                        file_name="metrics.csv",
+                        mime="text/csv"
+                    )
+            
         except Exception as e:
             st.error(f"Error processing image: {e}")
-    
-    # Footer with additional information
-    st.markdown("---")
-    st.markdown("### About the Model")
-    st.markdown("""
-        This segmentation model is based on the SegFormer architecture and was fine-tuned on the Oxford-IIIT Pet dataset.
-        
-        **Key Performance Metrics:**
-        - Mean IoU (Intersection over Union): Measures overlap between predictions and ground truth
-        - Dice Coefficient: Similar to F1-score, balances precision and recall
-        
-        The model segments pet images into three semantic classes (background, border, and pet/foreground),
-        making it useful for applications like pet image editing, background removal, and object detection.
-    """)
+            if debug_mode:
+                import traceback
+                st.code(traceback.format_exc())
+
 
 if __name__ == "__main__":
+    # Try to configure GPU memory growth
+    try:
+        gpus = tf.config.experimental.list_physical_devices('GPU')
+        if gpus:
+            for gpu in gpus:
+                tf.config.experimental.set_memory_growth(gpu, True)
+    except Exception as e:
+        print(f"GPU configuration error: {e}")
+    
     main()