Prepare for HF Spaces: Fix paths, add .gitignore, update app logic

.gitignore

@@ -0,0 +1,24 @@
+# Models
+models/
+*.pth
+*.h5
+*.npy
+
+# Python
+__pycache__/
+*.py[cod]
+*$py.class
+
+# Environment
+.env
+.venv
+venv/
+ENV/
+
+# IDE
+.vscode/
+.idea/
+
+# OS
+.DS_Store
+Thumbs.db

src/app.py

@@ -1,206 +0,0 @@
-import streamlit as st
-import h5py
-import torch
-import numpy as np
-import matplotlib.pyplot as plt
-import yaml
-import os
-
-# Import models
-from mobilenetv2_model import LandslideModel as MobileNetV2Model
-from vgg16_model import LandslideModel as VGG16Model
-from resnet34_model import LandslideModel as ResNet34Model
-from efficientnetb0_model import LandslideModel as EfficientNetB0Model
-from mitb1_model import LandslideModel as MiTB1Model
-from inceptionv4_model import LandslideModel as InceptionV4Model
-from densenet121_model import LandslideModel as DenseNet121Model
-from deeplabv3plus_model import LandslideModel as DeepLabV3PlusModel
-from resnext50_32x4d_model import LandslideModel as ResNeXt50_32X4DModel
-from se_resnet50_model import LandslideModel as SEResNet50Model
-from se_resnext50_32x4d_model import LandslideModel as SEResNeXt50_32X4DModel
-from segformer_model import LandslideModel as SegFormerB2Model
-from inceptionresnetv2_model import LandslideModel as InceptionResNetV2Model
-
-# Load the configuration file
-config = """
-model_config:
-  model_type: "mobilenet_v2"
-  in_channels: 14
-  num_classes: 1
-  encoder_weights: "imagenet"
-  wce_weight: 0.5
-
-dataset_config:
-  num_classes: 1
-  num_channels: 14
-  channels: [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14]
-  normalize: False
-
-train_config:
-  dataset_path: ""
-  checkpoint_path: "checkpoints"
-  seed: 42
-  train_val_split: 0.8
-  batch_size: 16
-  num_epochs: 100
-  lr: 0.001
-  device: "cuda:0"
-  save_config: True
-  experiment_name: "mobilenet_v2"
-
-logging_config:
-  wandb_project: "l4s"
-  wandb_entity: "Silvamillion"
-"""
-
-config = yaml.safe_load(config)
-
-# Model descriptions
-model_descriptions = {
-    "MobileNetV2": {"path": "mobilenetv2.pth", "type": "mobilenet_v2", "description": "MobileNetV2 is a lightweight deep learning model for image classification and segmentation."},
-    "VGG16": {"path": "vgg16.pth", "type": "vgg16", "description": "VGG16 is a popular deep learning model known for its simplicity and depth."},
-    "ResNet34": {"path": "resnet34.pth", "type": "resnet34", "description": "ResNet34 is a deep residual network that helps in training very deep networks."},
-    "EfficientNetB0": {"path": "effucientnetb0.pth", "type": "efficientnet_b0", "description": "EfficientNetB0 is part of the EfficientNet family, known for its efficiency and performance."},
-    "MiT-B1": {"path": "mitb1.pth", "type": "mit_b1", "description": "MiT-B1 is a transformer-based model designed for segmentation tasks."},
-    "InceptionV4": {"path": "inceptionv4.pth", "type": "inceptionv4", "description": "InceptionV4 is a convolutional neural network known for its inception modules."},
-    "DeepLabV3+": {"path": "deeplabv3.pth", "type": "deeplabv3+", "description": "DeepLabV3+ is an advanced model for semantic image segmentation."},
-    "DenseNet121": {"path": "densenet121.pth", "type": "densenet121", "description": "DenseNet121 is a densely connected convolutional network for image classification and segmentation."},
-    "ResNeXt50_32X4D": {"path": "resnext50-32x4d.pth", "type": "resnext50_32x4d", "description": "ResNeXt50_32X4D is a highly modularized network aimed at improving accuracy."},
-    "SEResNet50": {"path": "se_resnet50.pth", "type": "se_resnet50", "description": "SEResNet50 is a ResNet model with squeeze-and-excitation blocks for better feature recalibration."},
-    "SEResNeXt50_32X4D": {"path": "se_resnext50_32x4d.pth", "type": "se_resnext50_32x4d", "description": "SEResNeXt50_32X4D combines ResNeXt and SE blocks for improved performance."},
-    "SegFormerB2": {"path": "segformer.pth", "type": "segformer_b2", "description": "SegFormerB2 is a transformer-based model for semantic segmentation."},
-    "InceptionResNetV2": {"path": "inceptionresnetv2.pth", "type": "inceptionresnetv2", "description": "InceptionResNetV2 is a hybrid model combining Inception and ResNet architectures."},
-}
-
-# Streamlit app
-st.set_page_config(page_title="Landslide Detection", layout="wide")
-
-st.title("Landslide Detection")
-st.markdown("""
-## Instructions
-1. Select a model from the sidebar or choose to run all models.
-2. Upload one or more `.h5` files.
-3. The app will process the files and display the input image, prediction, and overlay.
-4. You can download the prediction results.
-""")
-
-# Sidebar for model selection
-st.sidebar.title("Model Selection")
-model_option = st.sidebar.radio("Choose an option", ["Select a single model", "Run all models"])
-if model_option == "Select a single model":
-    model_type = st.sidebar.selectbox("Select Model", list(model_descriptions.keys()))
-    config['model_config']['model_type'] = model_descriptions[model_type]['type']
-    if model_type == "DeepLabV3+":
-        model_class = DeepLabV3PlusModel
-    else:
-        model_class = locals()[model_type.replace("-", "") + "Model"]
-    model_path = model_descriptions[model_type]['path']
-
-    # Display model details in the sidebar
-    st.sidebar.markdown(f"**Model Type:** {model_descriptions[model_type]['type']}")
-    st.sidebar.markdown(f"**Model Path:** {model_descriptions[model_type]['path']}")
-    st.sidebar.markdown(f"**Description:** {model_descriptions[model_type]['description']}")
-
-# Main content
-st.header("Upload Data")
-uploaded_files = st.file_uploader("Choose .h5 files...", type="h5", accept_multiple_files=True)
-if uploaded_files:
-    for uploaded_file in uploaded_files:
-        st.write(f"Processing file: {uploaded_file.name}")
-        with st.spinner('Classifying...'):
-            with h5py.File(uploaded_file, 'r') as hdf:
-                data = np.array(hdf.get('img'))
-                data[np.isnan(data)] = 0.000001
-                channels = config["dataset_config"]["channels"]
-                image = np.zeros((128, 128, len(channels)))
-                for i, channel in enumerate(channels):
-                    image[:, :, i] = data[:, :, channel-1]
-
-            # Transform the image to the required format
-            image = image.transpose((2, 0, 1))  # (H, W, C) to (C, H, W)
-            image = torch.from_numpy(image).float().unsqueeze(0)  # Add batch dimension
-
-            if model_option == "Select a single model":
-                # Process the image with the selected model
-                st.write(f"Using model: {model_type}")
-
-                # Load the model
-                model = model_class(config)
-                model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu')), strict=False)
-                model.eval()
-
-                # Make prediction
-                with torch.no_grad():
-                    prediction = model(image)
-                    prediction = torch.sigmoid(prediction).cpu().numpy()
-
-                # Display prediction
-                st.header(f"Prediction Results - {model_type}")
-                fig, ax = plt.subplots(1, 3, figsize=(15, 5))
-                img = image.squeeze().permute(1, 2, 0).numpy()
-                img = (img - img.min()) / (img.max() - img.min())  # Normalize the image to [0, 1] range for display
-                ax[0].imshow(img[:, :, 1:4])  # Display first three channels as RGB
-                ax[0].set_title("Input Image")
-                ax[1].imshow(prediction.squeeze() > 0.5, cmap='plasma')  # Apply threshold
-                ax[1].set_title("Prediction")
-                ax[2].imshow(img[:, :, :3])  # Display first three channels as RGB
-                ax[2].imshow(prediction.squeeze() > 0.5, cmap='plasma', alpha=0.3)  # Overlay prediction
-                ax[2].set_title("Overlay")
-                st.pyplot(fig)
-
-                # Option to download the prediction
-                st.write(f"Download the prediction as a .npy file for {model_type}:")
-                npy_data = prediction.squeeze()
-                st.download_button(
-                    label=f"Download Prediction - {model_type}",
-                    data=npy_data.tobytes(),
-                    file_name=f"{uploaded_file.name.split('.')[0]}_{model_type}_prediction.npy",
-                    mime="application/octet-stream"
-                )
-
-            else:
-                # Process the image with each model
-                for model_name, model_info in model_descriptions.items():
-                    st.write(f"Using model: {model_name}")
-                    if model_name == "DeepLabV3+":
-                        model_class = DeepLabV3PlusModel
-                    else:
-                        model_class = locals()[model_name.replace("-", "") + "Model"]
-                    model_path = model_info['path']
-                    config['model_config']['model_type'] = model_info['type']
-
-                    # Load the model
-                    model = model_class(config)
-                    model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu')), strict=False)
-                    model.eval()
-
-                    # Make prediction
-                    with torch.no_grad():
-                        prediction = model(image)
-                        prediction = torch.sigmoid(prediction).cpu().numpy()
-
-                    # Display prediction
-                    st.header(f"Prediction Results - {model_name}")
-                    fig, ax = plt.subplots(1, 3, figsize=(15, 5))
-                    img = image.squeeze().permute(1, 2, 0).numpy()
-                    img = (img - img.min()) / (img.max() - img.min())  # Normalize the image to [0, 1] range for display
-                    ax[0].imshow(img[:, :, :3])  # Display first three channels as RGB
-                    ax[0].set_title("Input Image")
-                    ax[1].imshow(prediction.squeeze() > 0.5, cmap='plasma')  # Apply threshold
-                    ax[1].set_title("Prediction")
-                    ax[2].imshow(img[:, :, :3])  # Display first three channels as RGB
-                    ax[2].imshow(prediction.squeeze() > 0.5, cmap='plasma', alpha=0.3)  # Overlay prediction
-                    ax[2].set_title("Overlay")
-                    st.pyplot(fig)
-
-                    # Option to download the prediction
-                    st.write(f"Download the prediction as a .npy file for {model_name}:")
-                    npy_data = prediction.squeeze()
-                    st.download_button(
-                        label=f"Download Prediction - {model_name}",
-                        data=npy_data.tobytes(),
-                        file_name=f"{uploaded_file.name.split('.')[0]}_{model_name}_prediction.npy",
-                        mime="application/octet-stream"
-                    )
-
-st.success('Done!')

src/deeplabv3plus_model.py

@@ -10,7 +10,7 @@ from torch.optim.lr_scheduler import StepLR
 class smp_model(nn.Module):
     def __init__(self, in_channels, out_channels, model_type, num_classes, encoder_weights):
         super(smp_model, self).__init__()
-        if model_type == "deeplabv3+":
+        if model_type == "deeplabv3plus":
             self.model = smp.DeepLabV3Plus(
                 encoder_name="resnet50",  # Change this to a valid encoder
                 encoder_weights=encoder_weights,

src/densenet121_model.py

@@ -18,7 +18,7 @@ class smp_model(nn.Module):
         )
 
     def load_pretrained_weights(self):
-        state_dict = torch.load('/home/hks/MOU/DenseNet121_14C_L4S/densenet121-fbdb23505-trainWeights.pth', map_location='cpu')
+        # self.model.load_pretrained_weights()
         conv1_weight = state_dict['features.conv0.weight']
         new_conv1_weight = torch.zeros((conv1_weight.shape[0], 14, *conv1_weight.shape[2:]))
         new_conv1_weight[:, :3, :, :] = conv1_weight  # Copy weights for the first 3 channels
@@ -50,7 +50,7 @@ class LandslideModel(pl.LightningModule):
                                    model_type=model_type, 
                                    num_classes=num_classes, 
                                    encoder_weights=encoder_weights)
-            self.model.load_pretrained_weights()
+            # self.model.load_pretrained_weights()
 
         self.weights = torch.tensor([5], dtype=torch.float32).to(self.device)
         self.wce = nn.BCELoss(weight=self.weights)

src/download_all_models.py

@@ -0,0 +1,25 @@
+import sys
+import os
+
+# Add the parent directory to sys.path to allow imports from 'src'
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
+
+from src.model_downloader import ModelDownloader
+
+def download_all():
+    print("Starting download of all models...")
+    downloader = ModelDownloader()
+    models = downloader.list_available_models()
+    
+    for model_name in models:
+        try:
+            print(f"Checking/Downloading {model_name}...")
+            path = downloader.download_model(model_name)
+            print(f"✓ {model_name} is ready at {path}")
+        except Exception as e:
+            print(f"✗ Failed to download {model_name}: {e}")
+
+    print("\nAll downloads completed.")
+
+if __name__ == "__main__":
+    download_all()

src/inceptionresnetv2_model.py

@@ -11,7 +11,7 @@ class smp_model(nn.Module):
     def __init__(self, in_channels, out_channels, model_type, num_classes, encoder_weights):
         super(smp_model, self).__init__()
         self.model = smp.Unet(
-            encoder_name=model_type,
+            encoder_name="inceptionresnetv2",
             encoder_weights=encoder_weights,
             in_channels=in_channels,
             classes=num_classes,

src/model_downloader.py

@@ -8,11 +8,11 @@ from tqdm.auto import tqdm
 class ModelDownloader:
     def __init__(self):
         # Create models directory for caching
-        self.models_dir = Path("/app/models")
+        self.models_dir = Path("models").resolve()
         self.models_dir.mkdir(exist_ok=True)
         
         # HuggingFace model repository details
-        self.hf_model_url = "https://huggingface.co/harshinde/Sims/resolve/main/models/"
+        self.hf_model_url = "https://huggingface.co/harshinde/DeepSlide_Models/resolve/main/"
         
         # Model mapping with file names
         self.model_files = {
@@ -26,47 +26,47 @@ class ModelDownloader:
             },
             "efficientnetb0": {
                 "file": "efficientnetb0.pth",
-                "url": f"{self.hf_model_url}efficientnetb0.pth"
+                "url": f"{self.hf_model_url}effucientnetb0.pth"
             },
             "inceptionresnetv2": {
                 "file": "inceptionresnetv2.pth",
-                "id": "inceptionresnetv2"
+                "url": f"{self.hf_model_url}inceptionresnetv2.pth"
             },
             "inceptionv4": {
                 "file": "inceptionv4.pth",
-                "id": "inceptionv4"
+                "url": f"{self.hf_model_url}inceptionv4.pth"
             },
             "mitb1": {
                 "file": "mitb1.pth",
-                "id": "mitb1"
+                "url": f"{self.hf_model_url}mitb1.pth"
             },
             "mobilenetv2": {
                 "file": "mobilenetv2.pth",
-                "id": "mobilenetv2"
+                "url": f"{self.hf_model_url}mobilenetv2.pth"
             },
             "resnet34": {
                 "file": "resnet34.pth",
-                "id": "resnet34"
+                "url": f"{self.hf_model_url}resnet34.pth"
             },
             "resnext50_32x4d": {
                 "file": "resnext50-32x4d.pth",
-                "id": "resnext50-32x4d"
+                "url": f"{self.hf_model_url}resnext50-32x4d.pth"
             },
             "se_resnet50": {
                 "file": "se_resnet50.pth",
-                "id": "se_resnet50"
+                "url": f"{self.hf_model_url}se_resnet50.pth"
             },
             "se_resnext50_32x4d": {
                 "file": "se_resnext50_32x4d.pth",
-                "id": "se_resnext50_32x4d"
+                "url": f"{self.hf_model_url}se_resnext50_32x4d.pth"
             },
             "segformer": {
                 "file": "segformer.pth",
-                "id": "segformer"
+                "url": f"{self.hf_model_url}segformer.pth"
             },
             "vgg16": {
                 "file": "vgg16.pth",
-                "id": "vgg16"
+                "url": f"{self.hf_model_url}vgg16.pth"
             }
         }
 
@@ -86,7 +86,15 @@ class ModelDownloader:
 
         if not model_path.exists():
             print(f"Downloading {model_name} model...")
-            response = requests.get(model_info['url'], stream=True)
+            # Use 'url' if available, otherwise fallback or error (logic simplified for now as per plan)
+            if 'url' in model_info:
+                url = model_info['url']
+            else:
+                 # Fallback for models without explicit URL in the map (though all seem to have it or use ID)
+                 # Assuming the pattern from init for others
+                 url = f"{self.hf_model_url}{model_info['file']}"
+
+            response = requests.get(url, stream=True)
             response.raise_for_status()
             
             total_size = int(response.headers.get('content-length', 0))
@@ -102,55 +110,6 @@ class ModelDownloader:
             print(f"Model downloaded successfully to {model_path}")
         
         return str(model_path)
-        
-        # If model already exists, return path
-        if model_path.exists():
-            return str(model_path)
-
-        # Construct download URL for the specific model
-        download_url = f"{self.kaggle_model_url}/{model_info['id']}/1"
-        
-        try:
-            st.info(f"Downloading model {model_name} from Kaggle Models...")
-            progress_bar = st.progress(0)
-            
-            # Download with progress
-            response = requests.get(download_url, stream=True)
-            response.raise_for_status()
-            
-            total_size = int(response.headers.get('content-length', 0))
-            block_size = 1024  # 1 Kibibyte
-            
-            with open(model_path, 'wb') as f:
-                for i, data in enumerate(response.iter_content(block_size)):
-                    progress_bar.progress(min(i * block_size / total_size, 1.0))
-                    f.write(data)
-            
-            st.success(f"Successfully downloaded {model_name}")
-            return str(model_path)
-            
-        except requests.exceptions.RequestException as e:
-            raise Exception(f"Failed to download model from Kaggle: {str(e)}")
-
-    def get_model_path(self, model_name):
-        """
-        Get the path for a model file, downloading it from Kaggle if necessary
-        Args:
-            model_name (str): Name of the model (e.g., 'deeplabv3plus', 'densenet121', etc.)
-        Returns:
-            str: Path to the model file
-        """
-        if model_name not in self.model_files:
-            raise ValueError(f"Model {model_name} not found. Available models: {list(self.model_files.keys())}")
-        
-        model_info = self.model_files[model_name]
-        model_path = self.models_dir / model_info['file']
-        
-        # If model doesn't exist locally, download it
-        if not model_path.exists():
-            return self.download_from_kaggle(model_name)
-        
-        return str(model_path)
 
     def list_available_models(self):
         """

src/streamlit_app.py

@@ -1,10 +1,17 @@
 import streamlit as st
+import sys
+import os
+
+# Add the parent directory to sys.path to allow imports from 'src'
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
+
 import h5py
 import torch
 import numpy as np
 import matplotlib.pyplot as plt
 import yaml
 import os
+import io
 
 # Import models
 from src.mobilenetv2_model import LandslideModel as MobileNetV2Model
@@ -15,11 +22,12 @@ from src.mitb1_model import LandslideModel as MiTB1Model
 from src.inceptionv4_model import LandslideModel as InceptionV4Model
 from src.densenet121_model import LandslideModel as DenseNet121Model
 from src.deeplabv3plus_model import LandslideModel as DeepLabV3PlusModel
-from src.resnext50_32x4d_model import LandslideModel as ResNeXt50_32X4DModel
+from src.resnext50_32x4d_model import LandslideModel as ResNeXt50Model
 from src.se_resnet50_model import LandslideModel as SEResNet50Model
-from src.se_resnext50_32x4d_model import LandslideModel as SEResNeXt50_32X4DModel
+from src.se_resnext50_32x4d_model import LandslideModel as SEResNeXt50Model
 from src.segformer_model import LandslideModel as SegFormerB2Model
 from src.inceptionresnetv2_model import LandslideModel as InceptionResNetV2Model
+from src.model_downloader import ModelDownloader
 
 # Define available models
 AVAILABLE_MODELS = {
@@ -31,10 +39,10 @@ AVAILABLE_MODELS = {
     "inceptionv4": {"name": "InceptionV4", "type": "inception_v4"},
     "densenet121": {"name": "DenseNet121", "type": "densenet121"},
     "deeplabv3plus": {"name": "DeepLabV3Plus", "type": "deeplabv3plus"},
-    "resnext50": {"name": "ResNeXt50", "type": "resnext50_32x4d"},
-    "seresnet50": {"name": "SEResNet50", "type": "se_resnet50"},
-    "seresnext50": {"name": "SEResNeXt50", "type": "se_resnext50_32x4d"},
-    "segformerb2": {"name": "SegFormerB2", "type": "segformer_b2"},
+    "resnext50": {"name": "ResNeXt50", "type": "resnext50_32x4d", "downloader_key": "resnext50_32x4d"},
+    "seresnet50": {"name": "SEResNet50", "type": "se_resnet50", "downloader_key": "se_resnet50"},
+    "seresnext50": {"name": "SEResNeXt50", "type": "se_resnext50_32x4d", "downloader_key": "se_resnext50_32x4d"},
+    "segformerb2": {"name": "SegFormerB2", "type": "segformer_b2", "downloader_key": "segformer"},
     "inceptionresnetv2": {"name": "InceptionResNetV2", "type": "inception_resnet_v2"}
 }
 
@@ -43,13 +51,14 @@ MODEL_DESCRIPTIONS = {
     model_key: {
         "type": model_info["type"],
         "description": f"{model_info['name']} - A model for landslide detection and segmentation.",
-        "name": model_info["name"]
+        "name": model_info["name"],
+        "downloader_key": model_info.get("downloader_key", model_key)
     }
     for model_key, model_info in AVAILABLE_MODELS.items()
 }
 
 # Load the configuration file
-config = """
+config_str = """
 model_config:
   model_type: "mobilenet_v2"
   in_channels: 14
@@ -74,9 +83,89 @@ train_config:
   device: "cuda:0"
   save_config: True
   experiment_name: "mobilenet_v2"
+
+logging_config:
+  wandb_project: "l4s"
+  wandb_entity: "Silvamillion"
 """
 
-config = yaml.safe_load(config)
+config = yaml.safe_load(config_str)
+
+def process_and_visualize(model_key, model_info, image_tensor, original_image, uploaded_file_name):
+    """
+    Process the image with the selected model and visualize results.
+    """
+    try:
+        st.write(f"Using model: {model_info['name']}")
+        
+        # Update config for the specific model
+        current_config = config.copy()
+        current_config['model_config']['model_type'] = model_info['type']
+        
+        # Get the model class
+        model_class_name = AVAILABLE_MODELS[model_key]['name'].replace('-', '') + 'Model'
+        if model_class_name not in globals():
+             # Fallback for naming inconsistencies if any
+             # Try to find it in globals
+             pass
+        model_class = globals()[model_class_name]
+
+        # Initialize model downloader
+        downloader = ModelDownloader()
+        
+        # Download/get model path
+        download_key = model_info.get('downloader_key', model_key)
+        model_path = downloader.download_model(download_key)
+        st.info(f"Using model from: {model_path}")
+        
+        # Load the model
+        model = model_class(current_config)
+        model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu')), strict=False)
+        model.eval()
+
+        # Make prediction
+        with torch.no_grad():
+            prediction = model(image_tensor)
+            prediction = torch.sigmoid(prediction).cpu().numpy()
+
+        # Display prediction
+        st.header(f"Prediction Results - {model_info['name']}")
+        fig, ax = plt.subplots(1, 3, figsize=(15, 5))
+        
+        # Normalize image for display
+        img_display = original_image.transpose(1, 2, 0) # (C, H, W) -> (H, W, C)
+        img_display = (img_display - img_display.min()) / (img_display.max() - img_display.min())
+        
+        ax[0].imshow(img_display[:, :, :3])  # Display first three channels as RGB
+        ax[0].set_title("Input Image")
+        ax[0].axis('off')
+        
+        ax[1].imshow(prediction.squeeze(), cmap='plasma')  # Raw prediction map
+        ax[1].set_title("Prediction Probability")
+        ax[1].axis('off')
+        
+        ax[2].imshow(img_display[:, :, :3])
+        ax[2].imshow(prediction.squeeze() > 0.5, cmap='plasma', alpha=0.4)  # Overlay
+        ax[2].set_title("Overlay (Threshold > 0.5)")
+        ax[2].axis('off')
+        
+        st.pyplot(fig)
+        plt.close(fig)
+
+        # Download button
+        st.write(f"Download the prediction as a .npy file for {model_info['name']}:")
+        npy_data = prediction.squeeze()
+        st.download_button(
+            label=f"Download Prediction - {model_info['name']}",
+            data=npy_data.tobytes(),
+            file_name=f"{uploaded_file_name.split('.')[0]}_{model_key}_prediction.npy",
+            mime="application/octet-stream"
+        )
+        
+    except Exception as e:
+        st.error(f"Error with model {model_info['name']}: {str(e)}")
+        import traceback
+        st.error(traceback.format_exc())
 
 # Streamlit app
 st.set_page_config(page_title="Landslide Detection", layout="wide")
@@ -94,10 +183,10 @@ st.markdown("""
 st.sidebar.title("Model Selection")
 model_option = st.sidebar.radio("Choose an option", ["Select a single model", "Run all models"])
 
+selected_model_key = None
 if model_option == "Select a single model":
     selected_model_key = st.sidebar.selectbox("Select Model", list(MODEL_DESCRIPTIONS.keys()))
     selected_model_info = MODEL_DESCRIPTIONS[selected_model_key]
-    config['model_config']['model_type'] = selected_model_info['type']
     
     # Display model details in the sidebar
     st.sidebar.markdown("### Model Details")
@@ -127,7 +216,6 @@ if uploaded_files:
         with st.spinner('Classifying...'):
             try:
                 # Read the file directly using BytesIO
-                import io
                 bytes_data = uploaded_file.getvalue()
                 bytes_io = io.BytesIO(bytes_data)
                 
@@ -139,413 +227,52 @@ if uploaded_files:
                     data = np.array(hdf.get('img'))
                     data[np.isnan(data)] = 0.000001
                     channels = config["dataset_config"]["channels"]
+                    
+                    # Prepare image
+                    # Assuming data shape is (14, 128, 128) based on typical satellite data or (128, 128, 14)
+                    # The original code did: image[:, :, i] = data[band-1] implying data is (14, 128, 128) if accessed by index
+                    # But later it did data[:, :, channel-1] in the else block?
+                    # Let's check the original code logic again.
+                    # Original code had two different logic blocks for data loading!
+                    # Block 1 (single model): image[:, :, i] = data[band-1] -> implies data is (C, H, W)
+                    # Block 2 (all models): image[:, :, i] = data[:, :, channel-1] -> implies data is (H, W, C)
+                    
+                    # I will assume (C, H, W) is more standard for HDF5 'img' usually, but let's try to be robust or pick one.
+                    # Given the inconsistency, I'll check data shape.
+                    
                     image = np.zeros((128, 128, len(channels)))
+                    
+                    if data.ndim == 3:
+                        if data.shape[0] == 14: # (C, H, W)
+                             for i, band in enumerate(channels):
+                                image[:, :, i] = data[band-1, :, :]
+                        elif data.shape[2] == 14: # (H, W, C)
+                             for i, band in enumerate(channels):
+                                image[:, :, i] = data[:, :, band-1]
+                        else:
+                            st.warning(f"Unexpected data shape: {data.shape}. Assuming (C, H, W).")
+                            for i, band in enumerate(channels):
+                                if band-1 < data.shape[0]:
+                                     image[:, :, i] = data[band-1, :, :]
+                    else:
+                        st.error(f"Data has {data.ndim} dimensions, expected 3.")
+                        continue
 
-                    for i, band in enumerate(channels):
-                        image[:, :, i] = data[band-1]
-
-                    selected_channels = [image[:, :, i] for i in range(3)]
-                    image = np.transpose(image, (2, 0, 1))
+                    # Prepare for model (Batch, Channel, Height, Width)
+                    # image is currently (H, W, C)
+                    image_display = image.transpose(2, 0, 1) # (C, H, W)
+                    image_tensor = torch.from_numpy(image_display).unsqueeze(0).float() # (1, C, H, W)
 
                     if model_option == "Select a single model":
-                        # Get the model class from AVAILABLE_MODELS
-                        model_class_name = AVAILABLE_MODELS[selected_model_key]['name'].replace('-', '') + 'Model'
-                        model_class = locals()[model_class_name]
-
-                        # Initialize model downloader
-                        from model_downloader import ModelDownloader
-                        downloader = ModelDownloader()
-                        
-                        try:
-                            # Download/get model path
-                            model_path = downloader.download_model(selected_model_key)
-                            st.info(f"Using model from: {model_path}")
-                            
-                            # Load the model
-                            model = model_class(config)
-                            model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu')), strict=False)
-                            model.eval()
-
-                            # Make prediction
-                            with torch.no_grad():
-                                prediction = model(torch.from_numpy(image).unsqueeze(0).float())
-                                prediction = torch.sigmoid(prediction).numpy()
-
-                            st.header(f"Prediction Results - {selected_model_info['name']}")
-
-                            # Create columns for input image, prediction, and overlay
-                            col1, col2, col3 = st.columns(3)
-
-                            # Display input image
-                            with col1:
-                                st.write("Input Image")
-                                plt.figure(figsize=(8, 8))
-                                plt.imshow(selected_channels[0], cmap='viridis')
-                                plt.colorbar()
-                                plt.axis('off')
-                                st.pyplot(plt)
-
-                            # Display prediction
-                            with col2:
-                                st.write("Prediction")
-                                plt.figure(figsize=(8, 8))
-                                plt.imshow(prediction.squeeze(), cmap='viridis')
-                                plt.colorbar()
-                                plt.axis('off')
-                                st.pyplot(plt)
-
-                            # Display overlay
-                            with col3:
-                                st.write("Overlay")
-                                plt.figure(figsize=(8, 8))
-                                plt.imshow(selected_channels[0], cmap='viridis')
-                                plt.imshow(prediction.squeeze(), cmap='viridis', alpha=0.5)
-                                plt.colorbar()
-                                plt.axis('off')
-                                st.pyplot(plt)
-
-                            # Download button for prediction
-                            st.write(f"Download the prediction as a .npy file for {selected_model_info['name']}:")
-                            npy_data = prediction.squeeze()
-                            st.download_button(
-                                label=f"Download Prediction - {selected_model_info['name']}",
-                                data=npy_data.tobytes(),
-                                file_name=f"{uploaded_file.name.split('.')[0]}_{selected_model_key}_prediction.npy",
-                                mime="application/octet-stream"
-                            )
-
-                        except Exception as e:
-                            st.error(f"Error with model {selected_model_info['name']}: {str(e)}")
+                        process_and_visualize(selected_model_key, selected_model_info, image_tensor, image_display, uploaded_file.name)
                     else:
-                        # Process the image with each model
                         for model_key, model_info in MODEL_DESCRIPTIONS.items():
-                            st.write(f"Using model: {model_info['name']}")
-                            config['model_config']['model_type'] = model_info['type']
-                            
-                            # Get the model class from AVAILABLE_MODELS
-                            model_class_name = AVAILABLE_MODELS[model_key]['name'].replace('-', '') + 'Model'
-                            model_class = locals()[model_class_name]
-
-                            # Initialize model downloader
-                            from model_downloader import ModelDownloader
-                            downloader = ModelDownloader()
-                            
-                            try:
-                                # Download/get model path
-                                model_path = downloader.download_model(model_key)
-                                st.info(f"Using model from: {model_path}")
-                                
-                                # Load the model
-                                model = model_class(config)
-                                model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu')), strict=False)
-                                model.eval()
-
-                                # Make prediction
-                                with torch.no_grad():
-                                    prediction = model(torch.from_numpy(image).unsqueeze(0).float())
-                                    prediction = torch.sigmoid(prediction).numpy()
-
-                                st.header(f"Prediction Results - {model_info['name']}")
-
-                                # Create columns for input image, prediction, and overlay
-                                col1, col2, col3 = st.columns(3)
-
-                                # Display input image
-                                with col1:
-                                    st.write("Input Image")
-                                    plt.figure(figsize=(8, 8))
-                                    plt.imshow(selected_channels[0], cmap='viridis')
-                                    plt.colorbar()
-                                    plt.axis('off')
-                                    st.pyplot(plt)
-
-                                # Display prediction
-                                with col2:
-                                    st.write("Prediction")
-                                    plt.figure(figsize=(8, 8))
-                                    plt.imshow(prediction.squeeze(), cmap='viridis')
-                                    plt.colorbar()
-                                    plt.axis('off')
-                                    st.pyplot(plt)
-
-                                # Display overlay
-                                with col3:
-                                    st.write("Overlay")
-                                    plt.figure(figsize=(8, 8))
-                                    plt.imshow(selected_channels[0], cmap='viridis')
-                                    plt.imshow(prediction.squeeze(), cmap='viridis', alpha=0.5)
-                                    plt.colorbar()
-                                    plt.axis('off')
-                                    st.pyplot(plt)
-
-                                # Download button for prediction
-                                st.write(f"Download the prediction as a .npy file for {model_info['name']}:")
-                                npy_data = prediction.squeeze()
-                                st.download_button(
-                                    label=f"Download Prediction - {model_info['name']}",
-                                    data=npy_data.tobytes(),
-                                    file_name=f"{uploaded_file.name.split('.')[0]}_{model_key}_prediction.npy",
-                                    mime="application/octet-stream"
-                                )
-
-                            except Exception as e:
-                                st.error(f"Error with model {model_info['name']}: {str(e)}")
-                                continue
+                            process_and_visualize(model_key, model_info, image_tensor, image_display, uploaded_file.name)
 
             except Exception as e:
                 st.error(f"Error processing file {uploaded_file.name}: {str(e)}")
+                import traceback
+                st.error(traceback.format_exc())
                 continue
-import h5py
-import torch
-import numpy as np
-import matplotlib.pyplot as plt
-import yaml
-import os
-
-# Import models
-from src.mobilenetv2_model import LandslideModel as MobileNetV2Model
-from src.vgg16_model import LandslideModel as VGG16Model
-from src.resnet34_model import LandslideModel as ResNet34Model
-from src.efficientnetb0_model import LandslideModel as EfficientNetB0Model
-from src.mitb1_model import LandslideModel as MiTB1Model
-from src.inceptionv4_model import LandslideModel as InceptionV4Model
-from src.densenet121_model import LandslideModel as DenseNet121Model
-from src.deeplabv3plus_model import LandslideModel as DeepLabV3PlusModel
-from src.resnext50_32x4d_model import LandslideModel as ResNeXt50_32X4DModel
-from src.se_resnet50_model import LandslideModel as SEResNet50Model
-from src.se_resnext50_32x4d_model import LandslideModel as SEResNeXt50_32X4DModel
-from segformer_model import LandslideModel as SegFormerB2Model
-from inceptionresnetv2_model import LandslideModel as InceptionResNetV2Model
-
-# Define available models
-AVAILABLE_MODELS = {
-    "mobilenetv2": {"name": "MobileNetV2", "type": "mobilenet_v2"},
-    "vgg16": {"name": "VGG16", "type": "vgg16"},
-    "resnet34": {"name": "ResNet34", "type": "resnet34"},
-    "efficientnetb0": {"name": "EfficientNetB0", "type": "efficientnet_b0"},
-    "mitb1": {"name": "MiTB1", "type": "mitb1"},
-    "inceptionv4": {"name": "InceptionV4", "type": "inception_v4"},
-    "densenet121": {"name": "DenseNet121", "type": "densenet121"},
-    "deeplabv3plus": {"name": "DeepLabV3Plus", "type": "deeplabv3plus"},
-    "resnext50": {"name": "ResNeXt50", "type": "resnext50_32x4d"},
-    "seresnet50": {"name": "SEResNet50", "type": "se_resnet50"},
-    "seresnext50": {"name": "SEResNeXt50", "type": "se_resnext50_32x4d"},
-    "segformerb2": {"name": "SegFormerB2", "type": "segformer_b2"},
-    "inceptionresnetv2": {"name": "InceptionResNetV2", "type": "inception_resnet_v2"}
-}
-
-# Load the configuration file
-config = """
-model_config:
-  model_type: "mobilenet_v2"
-  in_channels: 14
-  num_classes: 1
-  encoder_weights: "imagenet"
-  wce_weight: 0.5
-
-dataset_config:
-  num_classes: 1
-  num_channels: 14
-  channels: [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14]
-  normalize: False
-
-train_config:
-  dataset_path: ""
-  checkpoint_path: "checkpoints"
-  seed: 42
-  train_val_split: 0.8
-  batch_size: 16
-  num_epochs: 100
-  lr: 0.001
-  device: "cuda:0"
-  save_config: True
-  experiment_name: "mobilenet_v2"
-
-logging_config:
-  wandb_project: "l4s"
-  wandb_entity: "Silvamillion"
-"""
-
-config = yaml.safe_load(config)
-
-# Model descriptions with their respective types and descriptions
-MODEL_DESCRIPTIONS = {
-    model_key: {
-        "type": model_info["type"],
-        "description": f"{model_info['name']} - A model for landslide detection and segmentation.",
-        "name": model_info["name"]
-    }
-    for model_key, model_info in AVAILABLE_MODELS.items()
-}
-
-# Streamlit app
-st.set_page_config(page_title="Landslide Detection", layout="wide")
-
-st.title("Landslide Detection")
-st.markdown("""
-## Instructions
-1. Select a model from the sidebar or choose to run all models.
-2. Upload one or more `.h5` files.
-3. The app will process the files and display the input image, prediction, and overlay.
-4. You can download the prediction results.
-""")
-
-# Sidebar for model selection
-st.sidebar.title("Model Selection")
-model_option = st.sidebar.radio("Choose an option", ["Select a single model", "Run all models"])
-if model_option == "Select a single model":
-    selected_model = st.sidebar.selectbox("Select Model", list(MODEL_DESCRIPTIONS.keys()))
-    config['model_config']['model_type'] = MODEL_DESCRIPTIONS[selected_model]['type']
-    
-    # Display model details in the sidebar
-    st.sidebar.markdown(f"**Model Name:** {MODEL_DESCRIPTIONS[selected_model]['name']}")
-    st.sidebar.markdown(f"**Model Type:** {MODEL_DESCRIPTIONS[selected_model]['type']}")
-    st.sidebar.markdown(f"**Description:** {MODEL_DESCRIPTIONS[selected_model]['description']}")
-
-# Main content
-st.header("Upload Data")
-
-# Initialize session state for error tracking if not exists
-if 'upload_errors' not in st.session_state:
-    st.session_state.upload_errors = []
-
-uploaded_files = st.file_uploader(
-    "Choose .h5 files...", 
-    type="h5", 
-    accept_multiple_files=True,
-    help="Upload your .h5 files here. Maximum file size is 200MB."
-)
-
-if uploaded_files:
-        for uploaded_file in uploaded_files:
-            st.write(f"Processing file: {uploaded_file.name}")
-            
-            # Display file details for debugging
-            st.write(f"File size: {uploaded_file.size} bytes")
-            
-            with st.spinner('Classifying...'):
-                try:
-                    # Read the file directly using BytesIO
-                    import io
-                    bytes_data = uploaded_file.getvalue()
-                    bytes_io = io.BytesIO(bytes_data)
-                    
-                    with h5py.File(bytes_io, 'r') as hdf:
-                        # Check if 'img' exists in the file
-                        if 'img' not in hdf:
-                            st.error(f"Error: 'img' dataset not found in {uploaded_file.name}")
-                            continue
-                            
-                        data = np.array(hdf.get('img'))
-                        data[np.isnan(data)] = 0.000001
-                        channels = config["dataset_config"]["channels"]
-                        image = np.zeros((128, 128, len(channels)))
-                        
-                except h5py.Error as he:
-                    st.error(f"H5PY Error processing {uploaded_file.name}: {str(he)}")
-                    continue
-                except Exception as e:
-                    st.error(f"Error processing {uploaded_file.name}: {str(e)}")
-                    continue
-                for i, channel in enumerate(channels):
-                    image[:, :, i] = data[:, :, channel-1]
-
-            # Transform the image to the required format
-            image = image.transpose((2, 0, 1))  # (H, W, C) to (C, H, W)
-            image = torch.from_numpy(image).float().unsqueeze(0)  # Add batch dimension
-
-            if model_option == "Select a single model":
-                # Process the image with the selected model
-                st.write(f"Using model: {model_type}")
-
-                # Load the model
-                model = model_class(config)
-                model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu')), strict=False)
-                model.eval()
-
-                # Make prediction
-                with torch.no_grad():
-                    prediction = model(image)
-                    prediction = torch.sigmoid(prediction).cpu().numpy()
-
-                # Display prediction
-                st.header(f"Prediction Results - {model_type}")
-                fig, ax = plt.subplots(1, 3, figsize=(15, 5))
-                img = image.squeeze().permute(1, 2, 0).numpy()
-                img = (img - img.min()) / (img.max() - img.min())  # Normalize the image to [0, 1] range for display
-                ax[0].imshow(img[:, :, 1:4])  # Display first three channels as RGB
-                ax[0].set_title("Input Image")
-                ax[1].imshow(prediction.squeeze() > 0.5, cmap='plasma')  # Apply threshold
-                ax[1].set_title("Prediction")
-                ax[2].imshow(img[:, :, :3])  # Display first three channels as RGB
-                ax[2].imshow(prediction.squeeze() > 0.5, cmap='plasma', alpha=0.3)  # Overlay prediction
-                ax[2].set_title("Overlay")
-                st.pyplot(fig)
-
-                # Option to download the prediction
-                st.write(f"Download the prediction as a .npy file for {model_type}:")
-                npy_data = prediction.squeeze()
-                st.download_button(
-                    label=f"Download Prediction - {model_type}",
-                    data=npy_data.tobytes(),
-                    file_name=f"{uploaded_file.name.split('.')[0]}_{model_type}_prediction.npy",
-                    mime="application/octet-stream"
-                )
-
-            else:
-                # Process the image with each model
-                for model_key, model_info in MODEL_DESCRIPTIONS.items():
-                    st.write(f"Using model: {model_info['name']}")
-                    config['model_config']['model_type'] = model_info['type']
-                    
-                    # Get the model class from AVAILABLE_MODELS
-                    model_class_name = AVAILABLE_MODELS[model_key]['name'].replace('-', '') + 'Model'
-                    model_class = locals()[model_class_name]
-
-                    # Initialize model downloader
-                    from model_downloader import ModelDownloader
-                    downloader = ModelDownloader()
-                    
-                    try:
-                        # Download/get model path
-                        model_path = downloader.download_model(model_name.lower())
-                        st.info(f"Using model from: {model_path}")
-                        
-                        # Load the model
-                        model = model_class(config)
-                        model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu')), strict=False)
-                        model.eval()
-                    except Exception as e:
-                        st.error(f"Error loading model {model_name}: {str(e)}")
-                        continue
-
-                    # Make prediction
-                    with torch.no_grad():
-                        prediction = model(image)
-                        prediction = torch.sigmoid(prediction).cpu().numpy()
-
-                    # Display prediction
-                    st.header(f"Prediction Results - {model_name}")
-                    fig, ax = plt.subplots(1, 3, figsize=(15, 5))
-                    img = image.squeeze().permute(1, 2, 0).numpy()
-                    img = (img - img.min()) / (img.max() - img.min())  # Normalize the image to [0, 1] range for display
-                    ax[0].imshow(img[:, :, :3])  # Display first three channels as RGB
-                    ax[0].set_title("Input Image")
-                    ax[1].imshow(prediction.squeeze() > 0.5, cmap='plasma')  # Apply threshold
-                    ax[1].set_title("Prediction")
-                    ax[2].imshow(img[:, :, :3])  # Display first three channels as RGB
-                    ax[2].imshow(prediction.squeeze() > 0.5, cmap='plasma', alpha=0.3)  # Overlay prediction
-                    ax[2].set_title("Overlay")
-                    st.pyplot(fig)
-
-                    # Option to download the prediction
-                    st.write(f"Download the prediction as a .npy file for {model_name}:")
-                    npy_data = prediction.squeeze()
-                    st.download_button(
-                        label=f"Download Prediction - {model_name}",
-                        data=npy_data.tobytes(),
-                        file_name=f"{uploaded_file.name.split('.')[0]}_{model_name}_prediction.npy",
-                        mime="application/octet-stream"
-                    )
 
 st.success('Done!')