Update app.py

app.py

@@ -124,12 +124,13 @@ def normalize_image(input_image):
     input_image = (input_image - MEAN) / tf.maximum(STD, backend.epsilon())
     return input_image
 
-def preprocess_image(image):
+def preprocess_image(image, from_dataset=False):
     """
-    Preprocess image for model input
+    Preprocess image for model input with special handling for dataset images
     
     Args:
         image: PIL Image to preprocess
+        from_dataset: Whether the image is from the original dataset
         
     Returns:
         Preprocessed image tensor, original image
@@ -140,23 +141,25 @@ def preprocess_image(image):
     # Store original image for display
     original_img = img_array.copy()
     
-    # Resize to target size with preserve_aspect_ratio=False
+    # Resize to target size
     img_resized = tf.image.resize(
         img_array, 
         (IMAGE_SIZE, IMAGE_SIZE),
         method='bilinear',
-        preserve_aspect_ratio=False,  # Ensure exact dimensions
+        preserve_aspect_ratio=False,
         antialias=True
     )
     
-    # Verify dimensions with assertion
-    tf.debugging.assert_equal(tf.shape(img_resized)[0:2], [IMAGE_SIZE, IMAGE_SIZE], 
-                             message="Image dimensions don't match expected size")
-    
-    # Normalize
-    img_normalized = normalize_image(img_resized)
+    # Special handling for dataset images
+    if from_dataset:
+        # The dataset already has specific dimensions, just normalize
+        # Skip additional preprocessing that might have been applied
+        img_normalized = normalize_image(img_resized)
+    else:
+        # Regular preprocessing for uploaded images
+        img_normalized = normalize_image(img_resized)
     
-    # Transpose from HWC to CHW (SegFormer expects channels first)
+    # Transpose from HWC to CHW (channels first)
     img_transposed = tf.transpose(img_normalized, (2, 0, 1))
     
     # Add batch dimension
@@ -329,6 +332,43 @@ def display_results_side_by_side(original_image, ground_truth_mask=None, predict
                 
             st.image(pred_display, use_column_width=True)
 
+def process_uploaded_mask(mask_array, from_dataset=False):
+    """
+    Process an uploaded mask to ensure it has the correct 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
+    """
+    # Check for RGBA format and convert to RGB if needed
+    if len(mask_array.shape) == 3 and mask_array.shape[2] == 4:
+        # Convert RGBA to RGB (discard alpha channel)
+        mask_array = mask_array[:,:,:3]
+    
+    # Convert RGB to grayscale if needed
+    if len(mask_array.shape) == 3 and mask_array.shape[2] >= 3:
+        # Convert RGB to grayscale
+        mask_array = cv2.cvtColor(mask_array, cv2.COLOR_RGB2GRAY)
+    
+    if from_dataset:
+        # For dataset masks (saved from your colab code):
+        # Create an empty mask with the same shape
+        processed_mask = np.zeros_like(mask_array)
+        
+        # Map the values correctly:
+        # Original dataset uses 1,2,3 but your app expects 0,1,2
+        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, we assume they have correct class values
+        return mask_array
+
 def main():
     st.title("🐶 Pet Segmentation with SegFormer")
     st.markdown("""
@@ -364,6 +404,9 @@ def main():
         step=0.1
     )
     
+    # Add this checkbox to your app's UI
+    dataset_image = st.sidebar.checkbox("Image is from the Oxford-IIIT Pet dataset")
+    
     # Load model
     with st.spinner("Loading SegFormer model..."):
         model = load_model()
@@ -395,7 +438,7 @@ def main():
             with st.spinner("Generating segmentation mask..."):
                 try:
                     # Preprocess the image
-                    img_tensor, original_img = preprocess_image(image)
+                    img_tensor, original_img = preprocess_image(image, from_dataset=dataset_image)
                     
                     # Print shape to debug
                     st.write(f"DEBUG - Input tensor shape: {img_tensor.shape}")
@@ -470,37 +513,48 @@ def main():
                     # 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)
+                    # Load the raw mask
+                    raw_mask = np.array(Image.open(mask_io))
+                    
+                    # Show debug info
+                    st.write(f"Debug - Raw mask shape: {raw_mask.shape}")
+                    st.write(f"Debug - Raw mask unique values: {np.unique(raw_mask)}")
+                    
+                    # Process the mask based on source
+                    processed_gt_mask = process_uploaded_mask(raw_mask, from_dataset=dataset_image)
+                    
+                    # Resize for IoU calculation
+                    gt_mask_resized = cv2.resize(processed_gt_mask, (OUTPUT_SIZE, OUTPUT_SIZE), 
+                                               interpolation=cv2.INTER_NEAREST)
+                    
+                    # Resize prediction for comparison
+                    pred_mask_resized = cv2.resize(mask, (OUTPUT_SIZE, OUTPUT_SIZE), 
+                                                 interpolation=cv2.INTER_NEAREST)
+                    
+                    # Show processed values
+                    st.write(f"Debug - Processed GT mask unique values: {np.unique(gt_mask_resized)}")
+                    st.write(f"Debug - Prediction mask unique values: {np.unique(pred_mask_resized)}")
                     
                     # 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)
+                    iou_score = calculate_iou(gt_mask_resized, pred_mask_resized)
                     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)
+                        bg_iou = calculate_iou(gt_mask_resized, pred_mask_resized, 0)
                         st.metric("Background IoU", f"{bg_iou:.4f}")
                     with col2:
-                        border_iou = calculate_iou(gt_mask, resized_mask, 1)
+                        border_iou = calculate_iou(gt_mask_resized, pred_mask_resized, 1)
                         st.metric("Border IoU", f"{border_iou:.4f}")
                     with col3:
-                        fg_iou = calculate_iou(gt_mask, resized_mask, 2)
+                        fg_iou = calculate_iou(gt_mask_resized, pred_mask_resized, 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))
+                    
+                    # For display, create a colorized version of the ground truth
+                    gt_mask_for_display = colorize_mask(processed_gt_mask)
                     
                     # Side-by-side display
                     display_results_side_by_side(
@@ -508,10 +562,11 @@ def main():
                         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.")
+                    import traceback
+                    st.code(traceback.format_exc())  # Show detailed error trace
                     
                     # Even with an error, try to display results without the ground truth
                     display_results_side_by_side(