add

app.py

@@ -27,33 +27,20 @@ STD = tf.constant([0.229, 0.224, 0.225])
 ID2LABEL = {0: "background", 1: "border", 2: "foreground/pet"}
 NUM_CLASSES = len(ID2LABEL)
 
+
 @st.cache_resource
 def download_model_from_drive():
-    """
-    Download model from Google Drive and return the local path.
-    """
-    model_dir = os.path.join("models", "saved_models")
-    os.makedirs(model_dir, exist_ok=True)
-    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):
+    model_path = "tf_model.h5"
+    if not os.path.exists(model_path):       
+        url = "https://drive.google.com/file/d/1XObpqG8qZ7YUyiRKbpVvxX11yQSK8Y_3/view?usp=sharing"  
         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!")
+            gdown.download(url, model_path, quiet=False)
+            st.success("Model downloaded successfully from Google Drive.")
         except Exception as e:
-            st.error(f"Error downloading model: {e}")
+            st.error(f"Failed to download model: {e}")
             return None
     else:
         st.info("Model already exists locally.")
-    
     return model_path
 
 @st.cache_resource