add

app.py

@@ -30,56 +30,29 @@ NUM_CLASSES = len(ID2LABEL)
 @st.cache_resource
 def download_model_from_drive():
     """
-    Download model from Google Drive
-    
-    Returns:
-        Path to downloaded model
+    Download model from Google Drive and return the local path.
     """
-    # Define paths
     model_dir = os.path.join("models", "saved_models")
     os.makedirs(model_dir, exist_ok=True)
-    model_path = os.path.join(model_dir, "t5_model.h5")
-    
-    # Check if model already exists
+    model_path = os.path.join(model_dir, "segformer_model")
+
+    # Google Drive file ID
+    file_id = "1XObpqG8qZ7YUyiRKbpVvxX11yQSK8Y_3"
+    url = f"https://drive.google.com/uc?id={file_id}"
+
     if not os.path.exists(model_path):
-        with st.spinner("Downloading model from Google Drive..."):
-            try:
-                # Direct download URL for t5_model.h5
-                url = "https://drive.google.com/uc?id=1XObpqG8qZ7YUyiRKbpVvxX11yQSK8Y_3"
-                
-                st.info(f"Downloading model from {url}")
-                
-                # Try direct download with gdown using direct URL
-                try:
-                    output = gdown.download(url, model_path, quiet=False)
-                    if output is None:
-                        raise Exception("gdown failed with direct URL")
-                    st.success(f"Model downloaded to {output}")
-                except Exception as e1:
-                    st.warning(f"First download attempt failed: {str(e1)}")
-                    
-                    # Try with just the file ID
-                    file_id = "1XObpqG8qZ7YUyiRKbpVvxX11yQSK8Y_3"
-                    st.info(f"Trying alternative download with file ID: {file_id}")
-                    output = gdown.download(f"https://drive.google.com/uc?id={file_id}", 
-                                           model_path, quiet=False)
-                    if output is None:
-                        raise Exception("gdown failed with file ID")
-                    st.success("Model downloaded successfully!")
-                
-                # Verify the file exists and has content
-                if os.path.exists(model_path) and os.path.getsize(model_path) > 0:
-                    st.success(f"Verified file at {model_path} ({os.path.getsize(model_path)/1024/1024:.2f} MB)")
-                else:
-                    st.error(f"Downloaded file is missing or empty: {model_path}")
-                    return None
-                
-            except Exception as e:
-                st.error(f"Error downloading model: {str(e)}")
-                st.info("Trying to use the HuggingFace model instead...")
-                return None
+        try:
+            # Use gdown to download the file
+            with st.spinner("Downloading model from Google Drive..."):
+                gdown.download(url, model_path, quiet=False)
+            if not os.path.exists(model_path) or os.path.getsize(model_path) == 0:
+                raise Exception("Downloaded file is empty or missing.")
+            st.success("Model downloaded successfully!")
+        except Exception as e:
+            st.error(f"Error downloading model: {e}")
+            return None
     else:
-        st.info(f"Model already exists at {model_path}")
+        st.info("Model already exists locally.")
     
     return model_path
 
@@ -192,36 +165,29 @@ def create_mask(pred_mask):
     Convert model prediction to displayable mask
     
     Args:
-        pred_mask: Prediction from model
+        pred_mask: Prediction logits from the model
         
     Returns:
-        Processed mask for visualization
+        Processed mask (2D array)
     """
-    # Get the class with highest probability (argmax along class dimension)
     pred_mask = tf.math.argmax(pred_mask, axis=1)
-    
-    # Add channel dimension
-    pred_mask = tf.expand_dims(pred_mask, -1)
-    
-    # Resize to original image size
-    pred_mask = tf.image.resize(
-        pred_mask, 
-        (IMAGE_SIZE, IMAGE_SIZE), 
-        method="nearest"
-    )
-    
-    return pred_mask[0]
+    pred_mask = tf.squeeze(pred_mask, axis=0)  # Remove batch dimension
+    return pred_mask.numpy()
 
 def colorize_mask(mask):
     """
     Apply colors to segmentation mask
     
     Args:
-        mask: Segmentation mask
+        mask: Segmentation mask (2D array)
         
     Returns:
-        Colorized mask
+        Colorized mask (3D RGB array)
     """
+    # Ensure the mask is 2D
+    if len(mask.shape) > 2:
+        mask = np.squeeze(mask, axis=-1)
+    
     # Define colors for each class (RGB)
     colors = [
         [0, 0, 0],      # Background (black)
@@ -233,8 +199,7 @@ def colorize_mask(mask):
     rgb_mask = np.zeros((mask.shape[0], mask.shape[1], 3), dtype=np.uint8)
     
     for i, color in enumerate(colors):
-        # Find pixels of this class and assign color
-        class_mask = np.where(mask == i, 1, 0).astype(np.uint8)
+        class_mask = (mask == i).astype(np.uint8)
         for c in range(3):
             rgb_mask[:, :, c] += class_mask * color[c]