add

app.py

@@ -10,7 +10,6 @@ import io
 import gdown
 from transformers import TFSegformerForSemanticSegmentation
 
-# Set page configuration
 st.set_page_config(
     page_title="Pet Segmentation with SegFormer",
     page_icon="🐶",
@@ -39,24 +38,48 @@ def download_model_from_drive():
     # Define paths
     model_dir = os.path.join("models", "saved_models")
     os.makedirs(model_dir, exist_ok=True)
-    model_path = os.path.join(model_dir, "segformer_model")
+    model_path = os.path.join(model_dir, "t5_model.h5")
     
     # Check if model already exists
     if not os.path.exists(model_path):
         with st.spinner("Downloading model from Google Drive..."):
             try:
-                # Google Drive file ID from the shared link
-                file_id = "1XObpqG8qZ7YUyiRKbpVvxX11yQSK8Y_3"
+                # 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
                 
-                # Download the model file
-                url = f"https://drive.google.com/uc?id={file_id}"
-                gdown.download(url, model_path, quiet=False)
-                st.success("Model downloaded successfully!")
             except Exception as e:
                 st.error(f"Error downloading model: {str(e)}")
+                st.info("Trying to use the HuggingFace model instead...")
                 return None
     else:
-        st.info("Model already exists locally.")
+        st.info(f"Model already exists at {model_path}")
     
     return model_path
 
@@ -83,8 +106,28 @@ def load_model():
                 ignore_mismatched_sizes=True
             )
         else:
-            # Load downloaded model
-            model = TFSegformerForSemanticSegmentation.from_pretrained(model_path)
+            # 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")
         
         return model
     except Exception as e: