Update app.py

app.py

@@ -1136,89 +1136,61 @@ if 'processed_images' not in st.session_state:
 tab1, tab2, tab3 = st.tabs(["📥 Load Model", " Process Single Image", "📁 Process Multiple Images"])
 
 # Add this code in tab1 before the model file uploader
+
+# Replace the model file uploader section in tab1 with this code
 with tab1:
     st.markdown("<h3 style='color: #a78bfa;'>Load Segmentation Model</h3>", unsafe_allow_html=True)
     
-    # Add model type selection
-    model_type = st.selectbox(
-        "Select model architecture",
-        ["U-Net", "DeepLabV3+", "SegNet"],
-        index=0,
-        help="Select the architecture of the model you're uploading"
-    )
+    # Get list of available models in the models/ directory
+    available_models = []
+    try:
+        for file in os.listdir('models'):
+            if file.endswith('.h5'):
+                available_models.append(file)
+    except FileNotFoundError:
+        st.error("Models directory not found. Please create a 'models' directory with .h5 model files.")
     
-    # Update the model type in the segmentation object
-    st.session_state.segmentation.model_type = model_type.lower().replace('-', '')
+    if not available_models:
+        st.warning("No model files found in the models directory. Please add .h5 model files to the 'models' folder.")
     
+    # Add model selection dropdown
     col1, col2 = st.columns([3, 1])
     
     with col1:
-        model_file = st.file_uploader("Upload model weights (.h5 file)", type=["h5"])
-
+        selected_model = st.selectbox(
+            "Select a pre-trained model",
+            available_models,
+            index=0 if available_models else None,
+            help="Select one of the available pre-trained models"
+        )
     
     with col2:
         st.markdown("<br>", unsafe_allow_html=True)
         if st.button("Load Model", key="load_model_btn"):
-            if model_file is not None:
+            if selected_model:
                 with st.spinner("Loading model..."):
-                    # Save the uploaded model to a temporary file
-                    with tempfile.NamedTemporaryFile(delete=False, suffix='.h5') as tmp:
-                        tmp.write(model_file.getvalue())
-                        model_path = tmp.name
-                    
                     try:
+                        # Load the selected model from the models directory
+                        model_path = os.path.join('models', selected_model)
+                        
+                        # Determine model type from filename (optional)
+                        model_type = "unet"  # Default
+                        if "deeplabv3" in selected_model.lower():
+                            model_type = "deeplabv3plus"
+                        elif "segnet" in selected_model.lower():
+                            model_type = "segnet"
+                        
+                        # Update model type in the segmentation object
+                        st.session_state.segmentation.model_type = model_type
+                        
                         # Load the model
                         st.session_state.segmentation.load_trained_model(model_path)
                         st.session_state.model_loaded = True
-                        st.success("Model loaded successfully!")
+                        st.success(f"Model '{selected_model}' loaded successfully!")
                     except Exception as e:
                         st.error(f"Error loading model: {str(e)}")
-                    finally:
-                        # Clean up the temporary file
-                        os.unlink(model_path)
             else:
-                st.error("Please upload a model file (.h5)")
-    
-
-    if st.session_state.model_loaded:
-        st.markdown("<div class='card'>", unsafe_allow_html=True)
-        st.markdown("<h4 style='color: #a78bfa;'>Model Information</h4>", unsafe_allow_html=True)
-        
-        col1, col2, col3 = st.columns(3)
-        with col1:
-            st.markdown("<div class='metric-card'>", unsafe_allow_html=True)
-            # Display the correct model architecture based on the detected model type
-            model_arch_map = {
-                'unet': "U-Net",
-                'deeplabv3plus': "DeepLabV3+",
-                'segnet': "SegNet"
-            }
-            model_arch = model_arch_map.get(st.session_state.segmentation.model_type, "Unknown")
-            st.markdown(f"<p class='metric-value'>{model_arch}</p>", unsafe_allow_html=True)
-            st.markdown("<p class='metric-label'>Architecture</p>", unsafe_allow_html=True)
-            st.markdown("</div>", unsafe_allow_html=True)
-
-        
-        with col2:
-            st.markdown("<div class='metric-card'>", unsafe_allow_html=True)
-            st.markdown("<p class='metric-value'>11</p>", unsafe_allow_html=True)
-            st.markdown("<p class='metric-label'>Land Cover Classes</p>", unsafe_allow_html=True)
-            st.markdown("</div>", unsafe_allow_html=True)
-        
-        with col3:
-            st.markdown("<div class='metric-card'>", unsafe_allow_html=True)
-            st.markdown("<p class='metric-value'>256 x 256</p>", unsafe_allow_html=True)
-            st.markdown("<p class='metric-label'>Input Size</p>", unsafe_allow_html=True)
-            st.markdown("</div>", unsafe_allow_html=True)
-        
-        # Display legend
-        st.markdown("<h4 style='color: #a78bfa; margin-top: 20px;'>Land Cover Classes</h4>", unsafe_allow_html=True)
-        legend_img = create_legend()
-        st.image(legend_img, use_column_width=True)
-        
-        st.markdown("</div>", unsafe_allow_html=True)
-    else:
-        st.info("Please load a model to continue.")
+                st.error("Please select a model from the dropdown")
 
 
 with tab2: