add

app.py

@@ -32,7 +32,8 @@ NUM_CLASSES = len(ID2LABEL)
 def download_model_from_drive():
     model_path = "tf_model.h5"
     if not os.path.exists(model_path):       
-        url = "https://drive.google.com/file/d/1XObpqG8qZ7YUyiRKbpVvxX11yQSK8Y_3/view?usp=sharing"  
+        # Fix the Google Drive link format - this is why the download is failing
+        url = "https://drive.google.com/uc?id=1XObpqG8qZ7YUyiRKbpVvxX11yQSK8Y_3"  
         try:
             gdown.download(url, model_path, quiet=False)
             st.success("Model downloaded successfully from Google Drive.")
@@ -55,7 +56,7 @@ def load_model():
         # Download the model first
         model_path = download_model_from_drive()
         
-        if model_path is None:
+        if model_path is None or not os.path.exists(model_path):
             st.warning("Using default pretrained model since download failed")
             # Fall back to pretrained model
             model = TFSegformerForSemanticSegmentation.from_pretrained(
@@ -66,28 +67,30 @@ def load_model():
                 ignore_mismatched_sizes=True
             )
         else:
-            # Check if this is a Keras .h5 model or a HuggingFace model directory
-            if model_path.endswith('.h5'):
-                st.info("Loading Keras H5 model...")
-                # For a Keras .h5 file, use tf.keras.models.load_model
-                try:
-                    model = tf.keras.models.load_model(model_path)
-                    st.success("Keras model loaded successfully")
-                except Exception as ke:
-                    st.error(f"Error loading Keras model: {str(ke)}")
-                    st.warning("Falling back to pretrained model")
-                    model = 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
-                    )
-            else:
-                # For a HuggingFace model directory
-                st.info("Loading HuggingFace model...")
-                model = TFSegformerForSemanticSegmentation.from_pretrained(model_path)
-                st.success("HuggingFace model loaded successfully")
+            # For a HuggingFace model saved with SavedModel format
+            st.info("Loading SegFormer model...")
+            try:
+                # First try loading as HuggingFace model
+                model = 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
+                )
+                # Then load weights from h5 file
+                model.load_weights(model_path)
+                st.success("Model weights loaded successfully")
+            except Exception as e:
+                st.error(f"Error loading model weights: {str(e)}")
+                st.warning("Falling back to pretrained model")
+                model = 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 model
     except Exception as e:
@@ -192,6 +195,38 @@ def colorize_mask(mask):
     
     return rgb_mask
 
+def calculate_iou(y_true, y_pred, class_idx=None):
+    """
+    Calculate IoU (Intersection over Union) for segmentation masks
+    
+    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)
+        
+    Returns:
+        IoU score
+    """
+    if class_idx is not None:
+        # Binary IoU for specific class
+        y_true_class = (y_true == class_idx).astype(np.float32)
+        y_pred_class = (y_pred == class_idx).astype(np.float32)
+        
+        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)
+    else:
+        # Mean IoU across all classes
+        class_ious = []
+        for idx in range(NUM_CLASSES):
+            class_iou = calculate_iou(y_true, y_pred, idx)
+            class_ious.append(class_iou)
+        
+        iou = np.mean(class_ious)
+    
+    return iou
+
 def create_overlay(image, mask, alpha=0.5):
     """
     Create an overlay of mask on original image
@@ -266,6 +301,7 @@ def main():
     # Image upload
     st.header("Upload an Image")
     uploaded_image = st.file_uploader("Upload a pet image:", type=["jpg", "jpeg", "png"])
+    uploaded_mask = st.file_uploader("Upload ground truth mask (optional):", type=["png", "jpg", "jpeg"])
     
     # Sample images option
     st.markdown("### Or use a sample image:")
@@ -342,6 +378,42 @@ 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 and process the mask file
+                mask_data = uploaded_mask.read()
+                st.write(f"Uploaded mask size: {len(mask_data)} bytes")
+                
+                # Open the mask from bytes
+                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
+                iou_score = calculate_iou(gt_mask, 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, mask, 0)
+                    st.metric("Background IoU", f"{bg_iou:.4f}")
+                with col2:
+                    border_iou = calculate_iou(gt_mask, mask, 1)
+                    st.metric("Border IoU", f"{border_iou:.4f}")
+                with col3:
+                    fg_iou = calculate_iou(gt_mask, 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)