add

app.py

@@ -30,10 +30,13 @@ NUM_CLASSES = len(ID2LABEL)
 
 @st.cache_resource
 def download_model_from_drive():
-    model_path = "tf_model.h5"
-    if not os.path.exists(model_path):       
-        # Fix the Google Drive link format - this is why the download is failing
-        url = "https://drive.google.com/file/d/1XObpqG8qZ7YUyiRKbpVvxX11yQSK8Y_3/view?usp=sharing"
+    # Create a models directory
+    os.makedirs("models", exist_ok=True)
+    model_path = "models/best_model"
+    
+    if not os.path.exists(model_path):
+        # Fixed Google Drive URL format for gdown
+        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.")
@@ -44,6 +47,7 @@ def download_model_from_drive():
         st.info("Model already exists locally.")
     return model_path
 
+
 @st.cache_resource
 def load_model():
     """
@@ -53,58 +57,44 @@ def load_model():
         Loaded model
     """
     try:
-        # Download the model first
+        # First create a base model with the correct architecture
+        base_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
+        )
+        
+        # Download the trained weights
         model_path = download_model_from_drive()
         
-        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(
-                "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 saved with SavedModel format
-            st.info("Loading SegFormer model...")
+        if model_path is not None and os.path.exists(model_path):
+            st.info(f"Loading weights from {model_path}...")
             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")
+                # 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 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
+                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
+            
     except Exception as e:
-        st.error(f"Error loading model: {str(e)}")
-        st.error("Falling back to pretrained model")
+        st.error(f"Error in load_model: {e}")
+        st.warning("Using default pretrained model")
         # Fall back to pretrained model as a last resort
-        model = TFSegformerForSemanticSegmentation.from_pretrained(
+        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 model
 
 def normalize_image(input_image):
     """
@@ -161,7 +151,7 @@ def create_mask(pred_mask):
         Processed mask (2D array)
     """
     pred_mask = tf.math.argmax(pred_mask, axis=1)
-    pred_mask = tf.squeeze(pred_mask, axis=0)  # Remove batch dimension
+    pred_mask = tf.squeeze(pred_mask)
     return pred_mask.numpy()
 
 def colorize_mask(mask):
@@ -176,7 +166,7 @@ def colorize_mask(mask):
     """
     # Ensure the mask is 2D
     if len(mask.shape) > 2:
-        mask = np.squeeze(mask, axis=-1)
+        mask = np.squeeze(mask)
     
     # Define colors for each class (RGB)
     colors = [
@@ -298,152 +288,135 @@ def main():
     else:
         st.sidebar.success("Model loaded successfully!")
     
-    # Image upload
+    # Image upload section
     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:")
-    sample_dir = "samples"
-    
-    # Check if sample directory exists and contains images
-    sample_files = []
-    if os.path.exists(sample_dir):
-        sample_files = [f for f in os.listdir(sample_dir) if f.endswith(('.jpg', '.jpeg', '.png'))]
-    
-    if sample_files:
-        selected_sample = st.selectbox("Select a sample image:", sample_files)
-        use_sample = st.button("Use this sample")
-        
-        if use_sample:
-            with open(os.path.join(sample_dir, selected_sample), "rb") as file:
-                image_bytes = file.read()
-                uploaded_image = io.BytesIO(image_bytes)
-                st.success(f"Using sample image: {selected_sample}")
-    
     # Process uploaded image
     if uploaded_image is not None:
-        # Display original image
-        image = Image.open(uploaded_image)
-        
-        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)
-            
-            # Make prediction
-            prediction = model(pixel_values=img_tensor, training=False)
-            logits = prediction.logits
+        try:
+            # Read the image
+            image_bytes = uploaded_image.read()
+            image = Image.open(io.BytesIO(image_bytes))
+                
+            col1, col2 = st.columns(2)
             
-            # Create visualization mask
-            mask = create_mask(logits).numpy()
+            with col1:
+                st.subheader("Original Image")
+                st.image(image, caption="Uploaded Image", use_column_width=True)
             
-            # Colorize the mask
-            colorized_mask = colorize_mask(mask)
+            # Preprocess and predict
+            with st.spinner("Generating segmentation mask..."):
+                # Preprocess the image
+                img_tensor, original_img = preprocess_image(image)
+                
+                # Make prediction
+                outputs = model(pixel_values=img_tensor, training=False)
+                logits = outputs.logits
+                
+                # Create visualization mask
+                mask = create_mask(logits)
+                
+                # Colorize the mask
+                colorized_mask = colorize_mask(mask)
+                
+                # Create overlay
+                overlay = create_overlay(original_img, colorized_mask, alpha=overlay_opacity)
             
-            # Create overlay
-            overlay = create_overlay(original_img, colorized_mask, alpha=overlay_opacity)
-        
-        # Display results
-        with col2:
-            st.subheader("Segmentation Result")
-            st.image(overlay, caption="Segmentation Overlay", use_column_width=True)
-        
-        # Display segmentation details
-        st.header("Segmentation Details")
-        col1, col2, col3 = st.columns(3)
-        
-        with col1:
-            st.subheader("Background")
-            st.markdown("Areas surrounding the pet")
-            mask_bg = np.where(mask == 0, 255, 0).astype(np.uint8)
-            st.image(mask_bg, caption="Background", use_column_width=True)
+            # Display results
+            with col2:
+                st.subheader("Segmentation Result")
+                st.image(overlay, caption="Segmentation Overlay", use_column_width=True)
             
-        with col2:
-            st.subheader("Border")
-            st.markdown("Boundary around the pet")
-            mask_border = np.where(mask == 1, 255, 0).astype(np.uint8)
-            st.image(mask_border, caption="Border", use_column_width=True)
+            # Display segmentation details
+            st.header("Segmentation Details")
+            col1, col2, col3 = st.columns(3)
             
-        with col3:
-            st.subheader("Foreground (Pet)")
-            st.markdown("The pet itself")
-            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)
+            with col1:
+                st.subheader("Background")
+                st.markdown("Areas surrounding the pet")
+                mask_bg = np.where(mask == 0, 255, 0).astype(np.uint8)
+                st.image(mask_bg, caption="Background", use_column_width=True)
                 
-                # Calculate and display IoU
-                iou_score = calculate_iou(gt_mask, mask)
-                st.success(f"Mean IoU: {iou_score:.4f}")
+            with col2:
+                st.subheader("Border")
+                st.markdown("Boundary around the pet")
+                mask_border = np.where(mask == 1, 255, 0).astype(np.uint8)
+                st.image(mask_border, caption="Border", use_column_width=True)
                 
-                # 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)
-        
-        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()
+            with col3:
+                st.subheader("Foreground (Pet)")
+                st.markdown("The pet itself")
+                mask_fg = np.where(mask == 2, 255, 0).astype(np.uint8)
+                st.image(mask_fg, caption="Foreground", use_column_width=True)
             
-            st.download_button(
-                label="Download Segmentation Mask",
-                data=mask_bytes,
-                file_name="pet_segmentation_mask.png",
-                mime="image/png"
-            )
-        
-        with col2:
-            # Convert overlay to PNG for download
-            overlay_img = Image.fromarray(overlay)
-            overlay_bytes = io.BytesIO()
-            overlay_img.save(overlay_bytes, format='PNG')
-            overlay_bytes = overlay_bytes.getvalue()
+            # 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.download_button(
-                label="Download Overlay Image",
-                data=overlay_bytes,
-                file_name="pet_segmentation_overlay.png",
-                mime="image/png"
-            )
+            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()
+                
+                st.download_button(
+                    label="Download Segmentation Mask",
+                    data=mask_bytes,
+                    file_name="pet_segmentation_mask.png",
+                    mime="image/png"
+                )
+            
+            with col2:
+                # Convert overlay to PNG for download
+                overlay_img = Image.fromarray(overlay)
+                overlay_bytes = io.BytesIO()
+                overlay_img.save(overlay_bytes, format='PNG')
+                overlay_bytes = overlay_bytes.getvalue()
+                
+                st.download_button(
+                    label="Download Overlay Image",
+                    data=overlay_bytes,
+                    file_name="pet_segmentation_overlay.png",
+                    mime="image/png"
+                )
+        except Exception as e:
+            st.error(f"Error processing image: {e}")
     
     # Footer with additional information
     st.markdown("---")