Upload 14 files

src/app.py

@@ -0,0 +1,206 @@
+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

@@ -0,0 +1,216 @@
+import torch
+import torch.nn as nn
+import segmentation_models_pytorch as smp
+from torchmetrics import F1Score, Precision, Recall, JaccardIndex
+import pytorch_lightning as pl
+import wandb
+from torch.optim import Adam
+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+":
+            self.model = smp.DeepLabV3Plus(
+                encoder_name="resnet50",  # Change this to a valid encoder
+                encoder_weights=encoder_weights,
+                in_channels=in_channels,
+                classes=num_classes
+            )
+        elif model_type == "unet":
+            self.model = smp.Unet(
+                encoder_name="resnet50",
+                encoder_weights=encoder_weights,
+                in_channels=in_channels,
+                classes=num_classes,
+            )
+        else:
+            raise ValueError(f"Model type {model_type} not supported!")
+
+    def forward(self, x):
+        return self.model(x)
+
+class LandslideModel(pl.LightningModule):
+    def __init__(self, config, alpha=0.5):
+        super(LandslideModel, self).__init__()
+
+        model_type = config['model_config']['model_type']
+        in_channels = config['model_config']['in_channels']
+        num_classes = config['model_config']['num_classes']
+        self.alpha = alpha
+        self.lr = config['train_config']['lr']
+
+        if model_type == 'unet':
+            self.model = UNet(in_channels=in_channels, out_channels=num_classes)
+        else:
+            encoder_weights = config['model_config']['encoder_weights']
+            self.model = smp_model(in_channels=in_channels, 
+                                   out_channels=num_classes, 
+                                   model_type=model_type, 
+                                   num_classes=num_classes, 
+                                   encoder_weights=encoder_weights)
+
+        self.weights = torch.tensor([5], dtype=torch.float32).to(self.device)
+        self.wce = nn.BCELoss(weight=self.weights)
+
+        self.train_f1 = F1Score(task='binary')
+        self.val_f1 = F1Score(task='binary')
+
+        self.train_precision = Precision(task='binary')
+        self.val_precision = Precision(task='binary')
+
+        self.train_recall = Recall(task='binary')
+        self.val_recall = Recall(task='binary')
+
+        self.train_iou = JaccardIndex(task='binary')
+        self.val_iou = JaccardIndex(task='binary')
+
+    def forward(self, x):
+        return self.model(x)
+
+    def training_step(self, batch, batch_idx):
+        x, y = batch
+        y_hat = torch.sigmoid(self(x))
+
+        wce_loss = self.wce(y_hat, y)
+        dice = dice_loss(y_hat, y)
+
+        combined_loss = (1 - self.alpha) * wce_loss + self.alpha * dice
+
+        precision = self.train_precision(y_hat, y)
+        recall = self.train_recall(y_hat, y)
+        iou = self.train_iou(y_hat, y)
+        loss_f1 = self.train_f1(y_hat, y)
+
+        self.log('train_precision', precision)
+        self.log('train_recall', recall)
+        self.log('train_wce', wce_loss)
+        self.log('train_dice', dice)
+        self.log('train_iou', iou)
+        self.log('train_f1', loss_f1)
+        self.log('train_loss', combined_loss)
+        return {'loss': combined_loss}
+
+    def validation_step(self, batch, batch_idx):
+        x, y = batch
+        y_hat = torch.sigmoid(self(x))
+
+        wce_loss = self.wce(y_hat, y)
+        dice = dice_loss(y_hat, y)
+
+        combined_loss = (1 - self.alpha) * wce_loss + self.alpha * dice
+
+        precision = self.val_precision(y_hat, y)
+        recall = self.val_recall(y_hat, y)
+        iou = self.val_iou(y_hat, y)
+        loss_f1 = self.val_f1(y_hat, y)
+
+        self.log('val_precision', precision)
+        self.log('val_recall', recall)
+        self.log('val_wce', wce_loss)
+        self.log('val_dice', dice)
+        self.log('val_iou', iou)
+        self.log('val_f1', loss_f1)
+        self.log('val_loss', combined_loss)
+
+        if self.current_epoch % 10 == 0:
+            x = (x - x.min()) / (x.max() - x.min())
+            x = x[:, 0:3]
+            x = x.permute(0, 2, 3, 1)
+            y_hat = (y_hat > 0.5).float()
+
+            class_labels = {0: "no landslide", 1: "landslide"}
+
+            self.logger.experiment.log({
+                "image": wandb.Image(x[0].cpu().detach().numpy(), masks={ 
+                    "predictions": {
+                        "mask_data": y_hat[0][0].cpu().detach().numpy(),
+                        "class_labels": class_labels
+                    },
+                    "ground_truth": {
+                        "mask_data": y[0][0].cpu().detach().numpy(),
+                        "class_labels": class_labels
+                    }
+                })
+            })
+        return {'val_loss': combined_loss}
+
+    def configure_optimizers(self):
+        optimizer = Adam(self.parameters(), lr=self.lr)
+        scheduler = StepLR(optimizer, step_size=30, gamma=0.1)
+        return [optimizer], [scheduler]
+
+class Block(nn.Module):
+    def __init__(self, inputs=3, middles=64, outs=64):
+        super().__init__()
+
+        self.conv1 = nn.Conv2d(inputs, middles, 3, 1, 1)
+        self.conv2 = nn.Conv2d(middles, outs, 3, 1, 1)
+        self.relu = nn.ReLU()
+        self.bn = nn.BatchNorm2d(outs)
+        self.pool = nn.MaxPool2d(2, 2)
+
+    def forward(self, x):
+        x = self.relu(self.conv1(x))
+        x = self.relu(self.bn(self.conv2(x)))
+        return self.pool(x), x
+
+class UNet(nn.Module):
+    def __init__(self, in_channels=3, out_channels=1):
+        super().__init__()
+
+        self.en1 = Block(in_channels, 64, 64)
+        self.en2 = Block(64, 128, 128)
+        self.en3 = Block(128, 256, 256)
+        self.en4 = Block(256, 512, 512)
+        self.en5 = Block(512, 1024, 512)
+
+        self.upsample4 = nn.ConvTranspose2d(512, 512, 2, stride=2)
+        self.de4 = Block(1024, 512, 256)
+
+        self.upsample3 = nn.ConvTranspose2d(256, 256, 2, stride=2)
+        self.de3 = Block(512, 256, 128)
+
+        self.upsample2 = nn.ConvTranspose2d(128, 128, 2, stride=2)
+        self.de2 = Block(256, 128, 64)
+
+        self.upsample1 = nn.ConvTranspose2d(64, 64, 2, stride=2)
+        self.de1 = Block(128, 64, 64)
+
+        self.conv_last = nn.Conv2d(64, out_channels, kernel_size=1, stride=1, padding=0)
+
+    def forward(self, x):
+        x, e1 = self.en1(x)
+        x, e2 = self.en2(x)
+        x, e3 = self.en3(x)
+        x, e4 = self.en4(x)
+        _, x = self.en5(x)
+
+        x = self.upsample4(x)
+        x = torch.cat([x, e4], dim=1)
+        _, x = self.de4(x)
+
+        x = self.upsample3(x)
+        x = torch.cat([x, e3], dim=1)
+        _, x = self.de3(x)
+
+        x = self.upsample2(x)
+        x = torch.cat([x, e2], dim=1)
+        _, x = self.de2(x)
+
+        x = self.upsample1(x)
+        x = torch.cat([x, e1], dim=1)
+        _, x = self.de1(x)
+
+        x = self.conv_last(x)
+
+        return x
+
+def dice_loss(y_hat, y):
+    smooth = 1e-6
+    y_hat = y_hat.view(-1)
+    y = y.view(-1)
+    intersection = (y_hat * y).sum()
+    union = y_hat.sum() + y.sum()
+    dice = (2 * intersection + smooth) / (union + smooth)
+    return 1 - dice

src/densenet121_model.py

@@ -0,0 +1,218 @@
+import torch
+import torch.nn as nn
+import segmentation_models_pytorch as smp
+from torchmetrics import F1Score, Precision, Recall, JaccardIndex
+import pytorch_lightning as pl
+import wandb
+from torch.optim import Adam
+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__()
+        self.model = smp.Unet(
+            encoder_name="densenet121",  # This will be "densenet121"
+            encoder_weights=None,  # Load weights manually
+            in_channels=in_channels,
+            classes=num_classes,
+        )
+
+    def load_pretrained_weights(self):
+        state_dict = torch.load('/home/hks/MOU/DenseNet121_14C_L4S/densenet121-fbdb23505-trainWeights.pth', map_location='cpu')
+        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
+        state_dict['features.conv0.weight'] = new_conv1_weight
+        model_dict = self.model.encoder.state_dict()
+        model_dict.update(state_dict)
+        self.model.encoder.load_state_dict(model_dict)
+
+    def forward(self, x):
+        x = self.model(x)
+        return x
+
+class LandslideModel(pl.LightningModule):
+    def __init__(self, config, alpha=0.5):
+        super(LandslideModel, self).__init__()
+
+        model_type = config['model_config']['model_type']
+        in_channels = config['model_config']['in_channels']
+        num_classes = config['model_config']['num_classes']
+        self.alpha = alpha
+        self.lr = config['train_config']['lr']
+
+        if model_type == 'unet':
+            self.model = UNet(in_channels=in_channels, out_channels=num_classes)
+        else:
+            encoder_weights = config['model_config']['encoder_weights']
+            self.model = smp_model(in_channels=in_channels, 
+                                   out_channels=num_classes, 
+                                   model_type=model_type, 
+                                   num_classes=num_classes, 
+                                   encoder_weights=encoder_weights)
+            self.model.load_pretrained_weights()
+
+        self.weights = torch.tensor([5], dtype=torch.float32).to(self.device)
+        self.wce = nn.BCELoss(weight=self.weights)
+
+        self.train_f1 = F1Score(task='binary')
+        self.val_f1 = F1Score(task='binary')
+
+        self.train_precision = Precision(task='binary')
+        self.val_precision = Precision(task='binary')
+
+        self.train_recall = Recall(task='binary')
+        self.val_recall = Recall(task='binary')
+
+        self.train_iou = JaccardIndex(task='binary')
+        self.val_iou = JaccardIndex(task='binary')
+
+    def forward(self, x):
+        return self.model(x)
+
+    def training_step(self, batch, batch_idx):
+        x, y = batch
+        y_hat = torch.sigmoid(self(x))
+
+        wce_loss = self.wce(y_hat, y)
+        dice = dice_loss(y_hat, y)
+
+        combined_loss = (1 - self.alpha) * wce_loss + self.alpha * dice
+
+        precision = self.train_precision(y_hat, y)
+        recall = self.train_recall(y_hat, y)
+        iou = self.train_iou(y_hat, y)
+        loss_f1 = self.train_f1(y_hat, y)
+
+        self.log('train_precision', precision)
+        self.log('train_recall', recall)
+        self.log('train_wce', wce_loss)
+        self.log('train_dice', dice)
+        self.log('train_iou', iou)
+        self.log('train_f1', loss_f1)
+        self.log('train_loss', combined_loss)
+        return {'loss': combined_loss}
+
+    def validation_step(self, batch, batch_idx):
+        x, y = batch
+        y_hat = torch.sigmoid(self(x))
+
+        wce_loss = self.wce(y_hat, y)
+        dice = dice_loss(y_hat, y)
+
+        combined_loss = (1 - self.alpha) * wce_loss + self.alpha * dice
+
+        precision = self.val_precision(y_hat, y)
+        recall = self.val_recall(y_hat, y)
+        iou = self.val_iou(y_hat, y)
+        loss_f1 = self.val_f1(y_hat, y)
+
+        self.log('val_precision', precision)
+        self.log('val_recall', recall)
+        self.log('val_wce', wce_loss)
+        self.log('val_dice', dice)
+        self.log('val_iou', iou)
+        self.log('val_f1', loss_f1)
+        self.log('val_loss', combined_loss)
+
+        if self.current_epoch % 10 == 0:
+            x = (x - x.min()) / (x.max() - x.min())
+            x = x[:, 0:3]
+            x = x.permute(0, 2, 3, 1)
+            y_hat = (y_hat > 0.5).float()
+
+            class_labels = {0: "no landslide", 1: "landslide"}
+
+            self.logger.experiment.log({
+                "image": wandb.Image(x[0].cpu().detach().numpy(), masks={ 
+                    "predictions": {
+                        "mask_data": y_hat[0][0].cpu().detach().numpy(),
+                        "class_labels": class_labels
+                    },
+                    "ground_truth": {
+                        "mask_data": y[0][0].cpu().detach().numpy(),
+                        "class_labels": class_labels
+                    }
+                })
+            })
+        return {'val_loss': combined_loss}
+
+    def configure_optimizers(self):
+        optimizer = Adam(self.parameters(), lr=self.lr)
+        scheduler = StepLR(optimizer, step_size=30, gamma=0.1)
+        return [optimizer], [scheduler]
+
+class Block(nn.Module):
+    def __init__(self, inputs=3, middles=64, outs=64):
+        super().__init__()
+
+        self.conv1 = nn.Conv2d(inputs, middles, 3, 1, 1)
+        self.conv2 = nn.Conv2d(middles, outs, 3, 1, 1)
+        self.relu = nn.ReLU()
+        self.bn = nn.BatchNorm2d(outs)
+        self.pool = nn.MaxPool2d(2, 2)
+
+    def forward(self, x):
+        x = self.relu(self.conv1(x))
+        x = self.relu(self.bn(self.conv2(x)))
+        return self.pool(x), x
+
+class UNet(nn.Module):
+    def __init__(self, in_channels=3, out_channels=1):
+        super().__init__()
+
+        self.en1 = Block(in_channels, 64, 64)
+        self.en2 = Block(64, 128, 128)
+        self.en3 = Block(128, 256, 256)
+        self.en4 = Block(256, 512, 512)
+        self.en5 = Block(512, 1024, 512)
+
+        self.upsample4 = nn.ConvTranspose2d(512, 512, 2, stride=2)
+        self.de4 = Block(1024, 512, 256)
+
+        self.upsample3 = nn.ConvTranspose2d(256, 256, 2, stride=2)
+        self.de3 = Block(512, 256, 128)
+
+        self.upsample2 = nn.ConvTranspose2d(128, 128, 2, stride=2)
+        self.de2 = Block(256, 128, 64)
+
+        self.upsample1 = nn.ConvTranspose2d(64, 64, 2, stride=2)
+        self.de1 = Block(128, 64, 64)
+
+        self.conv_last = nn.Conv2d(64, out_channels, kernel_size=1, stride=1, padding=0)
+
+    def forward(self, x):
+        x, e1 = self.en1(x)
+        x, e2 = self.en2(x)
+        x, e3 = self.en3(x)
+        x, e4 = self.en4(x)
+        _, x = self.en5(x)
+
+        x = self.upsample4(x)
+        x = torch.cat([x, e4], dim=1)
+        _, x = self.de4(x)
+
+        x = self.upsample3(x)
+        x = torch.cat([x, e3], dim=1)
+        _, x = self.de3(x)
+
+        x = self.upsample2(x)
+        x = torch.cat([x, e2], dim=1)
+        _, x = self.de2(x)
+
+        x = self.upsample1(x)
+        x = torch.cat([x, e1], dim=1)
+        _, x = self.de1(x)
+
+        x = self.conv_last(x)
+
+        return x
+
+def dice_loss(y_hat, y):
+    smooth = 1e-6
+    y_hat = y_hat.view(-1)
+    y = y.view(-1)
+    intersection = (y_hat * y).sum()
+    union = y_hat.sum() + y.sum()
+    dice = (2 * intersection + smooth) / (union + smooth)
+    return 1 - dice

src/efficientnetb0_model.py

@@ -0,0 +1,209 @@
+import torch
+import torch.nn as nn
+import segmentation_models_pytorch as smp
+from torchmetrics import F1Score, Precision, Recall, JaccardIndex
+import pytorch_lightning as pl
+import wandb
+from torch.optim import Adam
+from torch.optim.lr_scheduler import StepLR
+
+class smp_model_efficientnetb0(nn.Module):
+    def __init__(self, in_channels, out_channels, model_type, num_classes, encoder_weights):
+        super(smp_model_efficientnetb0, self).__init__()
+        
+        # Use EfficientNetB0 pre-trained model as encoder
+        self.model = smp.Unet(
+            encoder_name='efficientnet-b0',
+            encoder_weights=encoder_weights,
+            in_channels=in_channels,  # The number of input channels, which is 14
+            classes=num_classes,       # Output classes, which is 1
+        )
+
+    def forward(self, x):
+        return self.model(x)
+
+class LandslideModel(pl.LightningModule):
+    def __init__(self, config, alpha=0.5):
+        super(LandslideModel, self).__init__()
+
+        model_type = config['model_config']['model_type']
+        in_channels = config['model_config']['in_channels']
+        num_classes = config['model_config']['num_classes']
+        self.alpha = alpha  # Assign the alpha value to the class variable
+        self.lr = config['train_config']['lr']
+
+        if model_type == 'unet':
+            self.model = UNet(in_channels=in_channels, out_channels=num_classes)
+        else:
+            encoder_weights = config['model_config']['encoder_weights']
+            # Use the custom smp_model_efficientnetb0 instead of smp_model
+            self.model = smp_model_efficientnetb0(in_channels=in_channels, 
+                                                  out_channels=num_classes, 
+                                                  model_type=model_type, 
+                                                  num_classes=num_classes, 
+                                                  encoder_weights=encoder_weights)
+
+        self.weights = torch.tensor([5], dtype=torch.float32).to(self.device)
+        self.wce = nn.BCELoss(weight=self.weights)
+
+        self.train_f1 = F1Score(task='binary')
+        self.val_f1 = F1Score(task='binary')
+
+        self.train_precision = Precision(task='binary')
+        self.val_precision = Precision(task='binary')
+
+        self.train_recall = Recall(task='binary')
+        self.val_recall = Recall(task='binary')
+
+        self.train_iou = JaccardIndex(task='binary')
+        self.val_iou = JaccardIndex(task='binary')
+
+    def forward(self, x):
+        return self.model(x)
+
+    def training_step(self, batch, batch_idx):
+        x, y = batch
+        y_hat = torch.sigmoid(self(x))
+
+        wce_loss = self.wce(y_hat, y)
+        dice = dice_loss(y_hat, y)
+
+        combined_loss = (1 - self.alpha) * wce_loss + self.alpha * dice
+
+        precision = self.train_precision(y_hat, y)
+        recall = self.train_recall(y_hat, y)
+        iou = self.train_iou(y_hat, y)
+        loss_f1 = self.train_f1(y_hat, y)
+
+        self.log('train_precision', precision)
+        self.log('train_recall', recall)
+        self.log('train_wce', wce_loss)
+        self.log('train_dice', dice)
+        self.log('train_iou', iou)
+        self.log('train_f1', loss_f1)
+        self.log('train_loss', combined_loss)
+        return {'loss': combined_loss}
+
+    def validation_step(self, batch, batch_idx):
+        x, y = batch
+        y_hat = torch.sigmoid(self(x))
+
+        wce_loss = self.wce(y_hat, y)
+        dice = dice_loss(y_hat, y)
+
+        combined_loss = (1 - self.alpha) * wce_loss + self.alpha * dice
+
+        precision = self.val_precision(y_hat, y)
+        recall = self.val_recall(y_hat, y)
+        iou = self.val_iou(y_hat, y)
+        loss_f1 = self.val_f1(y_hat, y)
+
+        self.log('val_precision', precision)
+        self.log('val_recall', recall)
+        self.log('val_wce', wce_loss)
+        self.log('val_dice', dice)
+        self.log('val_iou', iou)
+        self.log('val_f1', loss_f1)
+        self.log('val_loss', combined_loss)
+
+        if self.current_epoch % 10 == 0:
+            x = (x - x.min()) / (x.max() - x.min())
+            x = x[:, 0:3]
+            x = x.permute(0, 2, 3, 1)
+            y_hat = (y_hat > 0.5).float()
+
+            class_labels = {0: "no landslide", 1: "landslide"}  # Define class_labels here
+
+            self.logger.experiment.log({
+                "image": wandb.Image(x[0].cpu().detach().numpy(), masks={ 
+                    "predictions": {
+                        "mask_data": y_hat[0][0].cpu().detach().numpy(),
+                        "class_labels": class_labels
+                    },
+                    "ground_truth": {
+                        "mask_data": y[0][0].cpu().detach().numpy(),
+                        "class_labels": class_labels
+                    }
+                })
+            })
+        return {'val_loss': combined_loss}
+
+    def configure_optimizers(self):
+        optimizer = Adam(self.parameters(), lr=self.lr)
+        scheduler = StepLR(optimizer, step_size=30, gamma=0.1)
+        return [optimizer], [scheduler]
+
+class Block(nn.Module):
+    def __init__(self, inputs=3, middles=64, outs=64):
+        super().__init__()
+
+        self.conv1 = nn.Conv2d(inputs, middles, 3, 1, 1)
+        self.conv2 = nn.Conv2d(middles, outs, 3, 1, 1)
+        self.relu = nn.ReLU()
+        self.bn = nn.BatchNorm2d(outs)
+        self.pool = nn.MaxPool2d(2, 2)
+
+    def forward(self, x):
+        x = self.relu(self.conv1(x))
+        x = self.relu(self.bn(self.conv2(x)))
+        return self.pool(x), x
+
+class UNet(nn.Module):
+    def __init__(self, in_channels=3, out_channels=1):
+        super().__init__()
+
+        self.en1 = Block(in_channels, 64, 64)
+        self.en2 = Block(64, 128, 128)
+        self.en3 = Block(128, 256, 256)
+        self.en4 = Block(256, 512, 512)
+        self.en5 = Block(512, 1024, 512)
+
+        self.upsample4 = nn.ConvTranspose2d(512, 512, 2, stride=2)
+        self.de4 = Block(1024, 512, 256)
+
+        self.upsample3 = nn.ConvTranspose2d(256, 256, 2, stride=2)
+        self.de3 = Block(512, 256, 128)
+
+        self.upsample2 = nn.ConvTranspose2d(128, 128, 2, stride=2)
+        self.de2 = Block(256, 128, 64)
+
+        self.upsample1 = nn.ConvTranspose2d(64, 64, 2, stride=2)
+        self.de1 = Block(128, 64, 64)
+
+        self.conv_last = nn.Conv2d(64, out_channels, kernel_size=1, stride=1, padding=0)
+
+    def forward(self, x):
+        x, e1 = self.en1(x)
+        x, e2 = self.en2(x)
+        x, e3 = self.en3(x)
+        x, e4 = self.en4(x)
+        _, x = self.en5(x)
+
+        x = self.upsample4(x)
+        x = torch.cat([x, e4], dim=1)
+        _, x = self.de4(x)
+
+        x = self.upsample3(x)
+        x = torch.cat([x, e3], dim=1)
+        _, x = self.de3(x)
+
+        x = self.upsample2(x)
+        x = torch.cat([x, e2], dim=1)
+        _, x = self.de2(x)
+
+        x = self.upsample1(x)
+        x = torch.cat([x, e1], dim=1)
+        _, x = self.de1(x)
+
+        x = self.conv_last(x)
+
+        return x
+
+def dice_loss(y_hat, y):
+    smooth = 1e-6
+    y_hat = y_hat.view(-1)
+    y = y.view(-1)
+    intersection = (y_hat * y).sum()
+    union = y_hat.sum() + y.sum()
+    dice = (2 * intersection + smooth) / (union + smooth)
+    return 1 - dice

src/inceptionresnetv2_model.py

@@ -0,0 +1,207 @@
+import torch
+import torch.nn as nn
+import segmentation_models_pytorch as smp
+from torchmetrics import F1Score, Precision, Recall, JaccardIndex
+import pytorch_lightning as pl
+import wandb
+from torch.optim import Adam
+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__()
+        self.model = smp.Unet(
+            encoder_name=model_type,
+            encoder_weights=encoder_weights,
+            in_channels=in_channels,
+            classes=num_classes,
+        )
+
+    def forward(self, x):
+        x = self.model(x)
+        return x
+
+class LandslideModel(pl.LightningModule):
+    def __init__(self, config, alpha=0.5):
+        super(LandslideModel, self).__init__()
+
+        model_type = config['model_config']['model_type']
+        in_channels = config['model_config']['in_channels']
+        num_classes = config['model_config']['num_classes']
+        self.alpha = alpha
+        self.lr = config['train_config']['lr']
+
+        if model_type == 'unet':
+            self.model = UNet(in_channels=in_channels, out_channels=num_classes)
+        else:
+            encoder_weights = config['model_config']['encoder_weights']
+            self.model = smp_model(in_channels=in_channels, 
+                                   out_channels=num_classes, 
+                                   model_type=model_type, 
+                                   num_classes=num_classes, 
+                                   encoder_weights=encoder_weights)
+
+        self.weights = torch.tensor([5], dtype=torch.float32).to(self.device)
+        self.wce = nn.BCELoss(weight=self.weights)
+
+        self.train_f1 = F1Score(task='binary')
+        self.val_f1 = F1Score(task='binary')
+
+        self.train_precision = Precision(task='binary')
+        self.val_precision = Precision(task='binary')
+
+        self.train_recall = Recall(task='binary')
+        self.val_recall = Recall(task='binary')
+
+        self.train_iou = JaccardIndex(task='binary')
+        self.val_iou = JaccardIndex(task='binary')
+
+    def forward(self, x):
+        return self.model(x)
+
+    def training_step(self, batch, batch_idx):
+        x, y = batch
+        y_hat = torch.sigmoid(self(x))
+
+        wce_loss = self.wce(y_hat, y)
+        dice = dice_loss(y_hat, y)
+
+        combined_loss = (1 - self.alpha) * wce_loss + self.alpha * dice
+
+        precision = self.train_precision(y_hat, y)
+        recall = self.train_recall(y_hat, y)
+        iou = self.train_iou(y_hat, y)
+        loss_f1 = self.train_f1(y_hat, y)
+
+        self.log('train_precision', precision)
+        self.log('train_recall', recall)
+        self.log('train_wce', wce_loss)
+        self.log('train_dice', dice)
+        self.log('train_iou', iou)
+        self.log('train_f1', loss_f1)
+        self.log('train_loss', combined_loss)
+        return {'loss': combined_loss}
+
+    def validation_step(self, batch, batch_idx):
+        x, y = batch
+        y_hat = torch.sigmoid(self(x))
+
+        wce_loss = self.wce(y_hat, y)
+        dice = dice_loss(y_hat, y)
+
+        combined_loss = (1 - self.alpha) * wce_loss + self.alpha * dice
+
+        precision = self.val_precision(y_hat, y)
+        recall = self.val_recall(y_hat, y)
+        iou = self.val_iou(y_hat, y)
+        loss_f1 = self.val_f1(y_hat, y)
+
+        self.log('val_precision', precision)
+        self.log('val_recall', recall)
+        self.log('val_wce', wce_loss)
+        self.log('val_dice', dice)
+        self.log('val_iou', iou)
+        self.log('val_f1', loss_f1)
+        self.log('val_loss', combined_loss)
+
+        if self.current_epoch % 10 == 0:
+            x = (x - x.min()) / (x.max() - x.min())
+            x = x[:, 0:3]
+            x = x.permute(0, 2, 3, 1)
+            y_hat = (y_hat > 0.5).float()
+
+            class_labels = {0: "no landslide", 1: "landslide"}
+
+            self.logger.experiment.log({
+                "image": wandb.Image(x[0].cpu().detach().numpy(), masks={ 
+                    "predictions": {
+                        "mask_data": y_hat[0][0].cpu().detach().numpy(),
+                        "class_labels": class_labels
+                    },
+                    "ground_truth": {
+                        "mask_data": y[0][0].cpu().detach().numpy(),
+                        "class_labels": class_labels
+                    }
+                })
+            })
+        return {'val_loss': combined_loss}
+
+    def configure_optimizers(self):
+        optimizer = Adam(self.parameters(), lr=self.lr)
+        scheduler = StepLR(optimizer, step_size=30, gamma=0.1)
+        return [optimizer], [scheduler]
+
+class Block(nn.Module):
+    def __init__(self, inputs=3, middles=64, outs=64):
+        super().__init__()
+
+        self.conv1 = nn.Conv2d(inputs, middles, 3, 1, 1)
+        self.conv2 = nn.Conv2d(middles, outs, 3, 1, 1)
+        self.relu = nn.ReLU()
+        self.bn = nn.BatchNorm2d(outs)
+        self.pool = nn.MaxPool2d(2, 2)
+
+    def forward(self, x):
+        x = self.relu(self.conv1(x))
+        x = self.relu(self.bn(self.conv2(x)))
+        return self.pool(x), x
+
+class UNet(nn.Module):
+    def __init__(self, in_channels=3, out_channels=1):
+        super().__init__()
+
+        self.en1 = Block(in_channels, 64, 64)
+        self.en2 = Block(64, 128, 128)
+        self.en3 = Block(128, 256, 256)
+        self.en4 = Block(256, 512, 512)
+        self.en5 = Block(512, 1024, 512)
+
+        self.upsample4 = nn.ConvTranspose2d(512, 512, 2, stride=2)
+        self.de4 = Block(1024, 512, 256)
+
+        self.upsample3 = nn.ConvTranspose2d(256, 256, 2, stride=2)
+        self.de3 = Block(512, 256, 128)
+
+        self.upsample2 = nn.ConvTranspose2d(128, 128, 2, stride=2)
+        self.de2 = Block(256, 128, 64)
+
+        self.upsample1 = nn.ConvTranspose2d(64, 64, 2, stride=2)
+        self.de1 = Block(128, 64, 64)
+
+        self.conv_last = nn.Conv2d(64, out_channels, kernel_size=1, stride=1, padding=0)
+
+    def forward(self, x):
+        x, e1 = self.en1(x)
+        x, e2 = self.en2(x)
+        x, e3 = self.en3(x)
+        x, e4 = self.en4(x)
+        _, x = self.en5(x)
+
+        x = self.upsample4(x)
+        x = torch.cat([x, e4], dim=1)
+        _, x = self.de4(x)
+
+        x = self.upsample3(x)
+        x = torch.cat([x, e3], dim=1)
+        _, x = self.de3(x)
+
+        x = self.upsample2(x)
+        x = torch.cat([x, e2], dim=1)
+        _, x = self.de2(x)
+
+        x = self.upsample1(x)
+        x = torch.cat([x, e1], dim=1)
+        _, x = self.de1(x)
+
+        x = self.conv_last(x)
+
+        return x
+
+def dice_loss(y_hat, y):
+    smooth = 1e-6
+    y_hat = y_hat.view(-1)
+    y = y.view(-1)
+    intersection = (y_hat * y).sum()
+    union = y_hat.sum() + y.sum()
+    dice = (2 * intersection + smooth) / (union + smooth)
+    return 1 - dice

src/inceptionv4_model.py

@@ -0,0 +1,207 @@
+import torch
+import torch.nn as nn
+import segmentation_models_pytorch as smp
+from torchmetrics import F1Score, Precision, Recall, JaccardIndex
+import pytorch_lightning as pl
+import wandb
+from torch.optim import Adam
+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__()
+        self.model = smp.Unet(
+            encoder_name="inceptionv4",  # Corrected to string "inceptionv4"
+            encoder_weights="imagenet",  # Corrected to string "imagenet"
+            in_channels=in_channels,  # Use the original in_channels
+            classes=num_classes,
+        )
+
+    def forward(self, x):
+        x = self.model(x)
+        return x
+
+class LandslideModel(pl.LightningModule):
+    def __init__(self, config, alpha=0.5):
+        super(LandslideModel, self).__init__()
+
+        model_type = config['model_config']['model_type']
+        in_channels = config['model_config']['in_channels']
+        num_classes = config['model_config']['num_classes']
+        self.alpha = alpha
+        self.lr = config['train_config']['lr']
+
+        if model_type == 'unet':
+            self.model = UNet(in_channels=in_channels, out_channels=num_classes)
+        else:
+            encoder_weights = config['model_config']['encoder_weights']
+            self.model = smp_model(in_channels=in_channels, 
+                                   out_channels=num_classes, 
+                                   model_type=model_type, 
+                                   num_classes=num_classes, 
+                                   encoder_weights=encoder_weights)
+
+        self.weights = torch.tensor([5], dtype=torch.float32).to(self.device)
+        self.wce = nn.BCELoss(weight=self.weights)
+
+        self.train_f1 = F1Score(task='binary')
+        self.val_f1 = F1Score(task='binary')
+
+        self.train_precision = Precision(task='binary')
+        self.val_precision = Precision(task='binary')
+
+        self.train_recall = Recall(task='binary')
+        self.val_recall = Recall(task='binary')
+
+        self.train_iou = JaccardIndex(task='binary')
+        self.val_iou = JaccardIndex(task='binary')
+
+    def forward(self, x):
+        return self.model(x)
+
+    def training_step(self, batch, batch_idx):
+        x, y = batch
+        y_hat = torch.sigmoid(self(x))
+
+        wce_loss = self.wce(y_hat, y)
+        dice = dice_loss(y_hat, y)
+
+        combined_loss = (1 - self.alpha) * wce_loss + self.alpha * dice
+
+        precision = self.train_precision(y_hat, y)
+        recall = self.train_recall(y_hat, y)
+        iou = self.train_iou(y_hat, y)
+        loss_f1 = self.train_f1(y_hat, y)
+
+        self.log('train_precision', precision)
+        self.log('train_recall', recall)
+        self.log('train_wce', wce_loss)
+        self.log('train_dice', dice)
+        self.log('train_iou', iou)
+        self.log('train_f1', loss_f1)
+        self.log('train_loss', combined_loss)
+        return {'loss': combined_loss}
+
+    def validation_step(self, batch, batch_idx):
+        x, y = batch
+        y_hat = torch.sigmoid(self(x))
+
+        wce_loss = self.wce(y_hat, y)
+        dice = dice_loss(y_hat, y)
+
+        combined_loss = (1 - self.alpha) * wce_loss + self.alpha * dice
+
+        precision = self.val_precision(y_hat, y)
+        recall = self.val_recall(y_hat, y)
+        iou = self.val_iou(y_hat, y)
+        loss_f1 = self.val_f1(y_hat, y)
+
+        self.log('val_precision', precision)
+        self.log('val_recall', recall)
+        self.log('val_wce', wce_loss)
+        self.log('val_dice', dice)
+        self.log('val_iou', iou)
+        self.log('val_f1', loss_f1)
+        self.log('val_loss', combined_loss)
+
+        if self.current_epoch % 10 == 0:
+            x = (x - x.min()) / (x.max() - x.min())
+            x = x[:, 0:3]
+            x = x.permute(0, 2, 3, 1)
+            y_hat = (y_hat > 0.5).float()
+
+            class_labels = {0: "no landslide", 1: "landslide"}
+
+            self.logger.experiment.log({
+                "image": wandb.Image(x[0].cpu().detach().numpy(), masks={ 
+                    "predictions": {
+                        "mask_data": y_hat[0][0].cpu().detach().numpy(),
+                        "class_labels": class_labels
+                    },
+                    "ground_truth": {
+                        "mask_data": y[0][0].cpu().detach().numpy(),
+                        "class_labels": class_labels
+                    }
+                })
+            })
+        return {'val_loss': combined_loss}
+
+    def configure_optimizers(self):
+        optimizer = Adam(self.parameters(), lr=self.lr)
+        scheduler = StepLR(optimizer, step_size=30, gamma=0.1)
+        return [optimizer], [scheduler]
+
+class Block(nn.Module):
+    def __init__(self, inputs=3, middles=64, outs=64):
+        super().__init__()
+
+        self.conv1 = nn.Conv2d(inputs, middles, 3, 1, 1)
+        self.conv2 = nn.Conv2d(middles, outs, 3, 1, 1)
+        self.relu = nn.ReLU()
+        self.bn = nn.BatchNorm2d(outs)
+        self.pool = nn.MaxPool2d(2, 2)
+
+    def forward(self, x):
+        x = self.relu(self.conv1(x))
+        x = self.relu(self.bn(self.conv2(x)))
+        return self.pool(x), x
+
+class UNet(nn.Module):
+    def __init__(self, in_channels=3, out_channels=1):
+        super().__init__()
+
+        self.en1 = Block(in_channels, 64, 64)
+        self.en2 = Block(64, 128, 128)
+        self.en3 = Block(128, 256, 256)
+        self.en4 = Block(256, 512, 512)
+        self.en5 = Block(512, 1024, 512)
+
+        self.upsample4 = nn.ConvTranspose2d(512, 512, 2, stride=2)
+        self.de4 = Block(1024, 512, 256)
+
+        self.upsample3 = nn.ConvTranspose2d(256, 256, 2, stride=2)
+        self.de3 = Block(512, 256, 128)
+
+        self.upsample2 = nn.ConvTranspose2d(128, 128, 2, stride=2)
+        self.de2 = Block(256, 128, 64)
+
+        self.upsample1 = nn.ConvTranspose2d(64, 64, 2, stride=2)
+        self.de1 = Block(128, 64, 64)
+
+        self.conv_last = nn.Conv2d(64, out_channels, kernel_size=1, stride=1, padding=0)
+
+    def forward(self, x):
+        x, e1 = self.en1(x)
+        x, e2 = self.en2(x)
+        x, e3 = self.en3(x)
+        x, e4 = self.en4(x)
+        _, x = self.en5(x)
+
+        x = self.upsample4(x)
+        x = torch.cat([x, e4], dim=1)
+        _, x = self.de4(x)
+
+        x = self.upsample3(x)
+        x = torch.cat([x, e3], dim=1)
+        _, x = self.de3(x)
+
+        x = self.upsample2(x)
+        x = torch.cat([x, e2], dim=1)
+        _, x = self.de2(x)
+
+        x = self.upsample1(x)
+        x = torch.cat([x, e1], dim=1)
+        _, x = self.de1(x)
+
+        x = self.conv_last(x)
+
+        return x
+
+def dice_loss(y_hat, y):
+    smooth = 1e-6
+    y_hat = y_hat.view(-1)
+    y = y.view(-1)
+    intersection = (y_hat * y).sum()
+    union = y_hat.sum() + y.sum()
+    dice = (2 * intersection + smooth) / (union + smooth)
+    return 1 - dice

src/mitb1_model.py

@@ -0,0 +1,209 @@
+import torch
+import torch.nn as nn
+import segmentation_models_pytorch as smp
+from torchmetrics import F1Score, Precision, Recall, JaccardIndex
+import pytorch_lightning as pl
+import wandb
+from torch.optim import Adam
+from torch.optim.lr_scheduler import StepLR
+
+class smp_model_mitb1(nn.Module):
+    def __init__(self, in_channels, out_channels, model_type, num_classes, encoder_weights):
+        super(smp_model_mitb1, self).__init__()
+        self.conv = nn.Conv2d(in_channels, 3, kernel_size=1)
+        self.model = smp.Unet(
+            encoder_name="mit_b1",  # Corrected to string "mit_b1"
+            encoder_weights="imagenet",  # Corrected to string "imagenet"
+            in_channels=3,  # Set in_channels to 3 for MiT-B1
+            classes=num_classes,
+        )
+
+    def forward(self, x):
+        x = self.conv(x)
+        x = self.model(x)
+        return x
+
+class LandslideModel(pl.LightningModule):
+    def __init__(self, config, alpha=0.5):
+        super(LandslideModel, self).__init__()
+
+        model_type = config['model_config']['model_type']
+        in_channels = config['model_config']['in_channels']
+        num_classes = config['model_config']['num_classes']
+        self.alpha = alpha
+        self.lr = config['train_config']['lr']
+
+        if model_type == 'unet':
+            self.model = UNet(in_channels=in_channels, out_channels=num_classes)
+        else:
+            encoder_weights = config['model_config']['encoder_weights']
+            self.model = smp_model_mitb1(in_channels=in_channels, 
+                                         out_channels=num_classes, 
+                                         model_type=model_type, 
+                                         num_classes=num_classes, 
+                                         encoder_weights=encoder_weights)
+
+        self.weights = torch.tensor([5], dtype=torch.float32).to(self.device)
+        self.wce = nn.BCELoss(weight=self.weights)
+
+        self.train_f1 = F1Score(task='binary')
+        self.val_f1 = F1Score(task='binary')
+
+        self.train_precision = Precision(task='binary')
+        self.val_precision = Precision(task='binary')
+
+        self.train_recall = Recall(task='binary')
+        self.val_recall = Recall(task='binary')
+
+        self.train_iou = JaccardIndex(task='binary')
+        self.val_iou = JaccardIndex(task='binary')
+
+    def forward(self, x):
+        return self.model(x)
+
+    def training_step(self, batch, batch_idx):
+        x, y = batch
+        y_hat = torch.sigmoid(self(x))
+
+        wce_loss = self.wce(y_hat, y)
+        dice = dice_loss(y_hat, y)
+
+        combined_loss = (1 - self.alpha) * wce_loss + self.alpha * dice
+
+        precision = self.train_precision(y_hat, y)
+        recall = self.train_recall(y_hat, y)
+        iou = self.train_iou(y_hat, y)
+        loss_f1 = self.train_f1(y_hat, y)
+
+        self.log('train_precision', precision)
+        self.log('train_recall', recall)
+        self.log('train_wce', wce_loss)
+        self.log('train_dice', dice)
+        self.log('train_iou', iou)
+        self.log('train_f1', loss_f1)
+        self.log('train_loss', combined_loss)
+        return {'loss': combined_loss}
+
+    def validation_step(self, batch, batch_idx):
+        x, y = batch
+        y_hat = torch.sigmoid(self(x))
+
+        wce_loss = self.wce(y_hat, y)
+        dice = dice_loss(y_hat, y)
+
+        combined_loss = (1 - self.alpha) * wce_loss + self.alpha * dice
+
+        precision = self.val_precision(y_hat, y)
+        recall = self.val_recall(y_hat, y)
+        iou = self.val_iou(y_hat, y)
+        loss_f1 = self.val_f1(y_hat, y)
+
+        self.log('val_precision', precision)
+        self.log('val_recall', recall)
+        self.log('val_wce', wce_loss)
+        self.log('val_dice', dice)
+        self.log('val_iou', iou)
+        self.log('val_f1', loss_f1)
+        self.log('val_loss', combined_loss)
+
+        if self.current_epoch % 10 == 0:
+            x = (x - x.min()) / (x.max() - x.min())
+            x = x[:, 0:3]
+            x = x.permute(0, 2, 3, 1)
+            y_hat = (y_hat > 0.5).float()
+
+            class_labels = {0: "no landslide", 1: "landslide"}
+
+            self.logger.experiment.log({
+                "image": wandb.Image(x[0].cpu().detach().numpy(), masks={ 
+                    "predictions": {
+                        "mask_data": y_hat[0][0].cpu().detach().numpy(),
+                        "class_labels": class_labels
+                    },
+                    "ground_truth": {
+                        "mask_data": y[0][0].cpu().detach().numpy(),
+                        "class_labels": class_labels
+                    }
+                })
+            })
+        return {'val_loss': combined_loss}
+
+    def configure_optimizers(self):
+        optimizer = Adam(self.parameters(), lr=self.lr)
+        scheduler = StepLR(optimizer, step_size=30, gamma=0.1)
+        return [optimizer], [scheduler]
+
+class Block(nn.Module):
+    def __init__(self, inputs=3, middles=64, outs=64):
+        super().__init__()
+
+        self.conv1 = nn.Conv2d(inputs, middles, 3, 1, 1)
+        self.conv2 = nn.Conv2d(middles, outs, 3, 1, 1)
+        self.relu = nn.ReLU()
+        self.bn = nn.BatchNorm2d(outs)
+        self.pool = nn.MaxPool2d(2, 2)
+
+    def forward(self, x):
+        x = self.relu(self.conv1(x))
+        x = self.relu(self.bn(self.conv2(x)))
+        return self.pool(x), x
+
+class UNet(nn.Module):
+    def __init__(self, in_channels=3, out_channels=1):
+        super().__init__()
+
+        self.en1 = Block(in_channels, 64, 64)
+        self.en2 = Block(64, 128, 128)
+        self.en3 = Block(128, 256, 256)
+        self.en4 = Block(256, 512, 512)
+        self.en5 = Block(512, 1024, 512)
+
+        self.upsample4 = nn.ConvTranspose2d(512, 512, 2, stride=2)
+        self.de4 = Block(1024, 512, 256)
+
+        self.upsample3 = nn.ConvTranspose2d(256, 256, 2, stride=2)
+        self.de3 = Block(512, 256, 128)
+
+        self.upsample2 = nn.ConvTranspose2d(128, 128, 2, stride=2)
+        self.de2 = Block(256, 128, 64)
+
+        self.upsample1 = nn.ConvTranspose2d(64, 64, 2, stride=2)
+        self.de1 = Block(128, 64, 64)
+
+        self.conv_last = nn.Conv2d(64, out_channels, kernel_size=1, stride=1, padding=0)
+
+    def forward(self, x):
+        x, e1 = self.en1(x)
+        x, e2 = self.en2(x)
+        x, e3 = self.en3(x)
+        x, e4 = self.en4(x)
+        _, x = self.en5(x)
+
+        x = self.upsample4(x)
+        x = torch.cat([x, e4], dim=1)
+        _, x = self.de4(x)
+
+        x = self.upsample3(x)
+        x = torch.cat([x, e3], dim=1)
+        _, x = self.de3(x)
+
+        x = self.upsample2(x)
+        x = torch.cat([x, e2], dim=1)
+        _, x = self.de2(x)
+
+        x = self.upsample1(x)
+        x = torch.cat([x, e1], dim=1)
+        _, x = self.de1(x)
+
+        x = self.conv_last(x)
+
+        return x
+
+def dice_loss(y_hat, y):
+    smooth = 1e-6
+    y_hat = y_hat.view(-1)
+    y = y.view(-1)
+    intersection = (y_hat * y).sum()
+    union = y_hat.sum() + y.sum()
+    dice = (2 * intersection + smooth) / (union + smooth)
+    return 1 - dice

src/mobilenetv2_model.py

@@ -0,0 +1,209 @@
+import torch
+import torch.nn as nn
+import segmentation_models_pytorch as smp
+from torchmetrics import F1Score, Precision, Recall, JaccardIndex
+import pytorch_lightning as pl
+import wandb
+from torch.optim import Adam
+from torch.optim.lr_scheduler import StepLR
+
+class smp_model_v2(nn.Module):
+    def __init__(self, in_channels, out_channels, model_type, num_classes, encoder_weights):
+        super(smp_model_v2, self).__init__()
+        
+        # Use MobileNetV2 pre-trained model as encoder
+        self.model = smp.Unet(
+            encoder_name='mobilenet_v2',
+            encoder_weights=encoder_weights,
+            in_channels=in_channels,  # The number of input channels, which is 14
+            classes=num_classes,       # Output classes, which is 1
+        )
+
+    def forward(self, x):
+        return self.model(x)
+
+class LandslideModel(pl.LightningModule):
+    def __init__(self, config, alpha=0.5):
+        super(LandslideModel, self).__init__()
+
+        model_type = config['model_config']['model_type']
+        in_channels = config['model_config']['in_channels']
+        num_classes = config['model_config']['num_classes']
+        self.alpha = alpha  # Assign the alpha value to the class variable
+        self.lr = config['train_config']['lr']
+
+        if model_type == 'unet':
+            self.model = UNet(in_channels=in_channels, out_channels=num_classes)
+        else:
+            encoder_weights = config['model_config']['encoder_weights']
+            # Use the custom smp_model_v2 instead of smp_model
+            self.model = smp_model_v2(in_channels=in_channels, 
+                                      out_channels=num_classes, 
+                                      model_type=model_type, 
+                                      num_classes=num_classes, 
+                                      encoder_weights=encoder_weights)
+
+        self.weights = torch.tensor([5], dtype=torch.float32).to(self.device)
+        self.wce = nn.BCELoss(weight=self.weights)
+
+        self.train_f1 = F1Score(task='binary')
+        self.val_f1 = F1Score(task='binary')
+
+        self.train_precision = Precision(task='binary')
+        self.val_precision = Precision(task='binary')
+
+        self.train_recall = Recall(task='binary')
+        self.val_recall = Recall(task='binary')
+
+        self.train_iou = JaccardIndex(task='binary')
+        self.val_iou = JaccardIndex(task='binary')
+
+    def forward(self, x):
+        return self.model(x)
+
+    def training_step(self, batch, batch_idx):
+        x, y = batch
+        y_hat = torch.sigmoid(self(x))
+
+        wce_loss = self.wce(y_hat, y)
+        dice = dice_loss(y_hat, y)
+
+        combined_loss = (1 - self.alpha) * wce_loss + self.alpha * dice
+
+        precision = self.train_precision(y_hat, y)
+        recall = self.train_recall(y_hat, y)
+        iou = self.train_iou(y_hat, y)
+        loss_f1 = self.train_f1(y_hat, y)
+
+        self.log('train_precision', precision)
+        self.log('train_recall', recall)
+        self.log('train_wce', wce_loss)
+        self.log('train_dice', dice)
+        self.log('train_iou', iou)
+        self.log('train_f1', loss_f1)
+        self.log('train_loss', combined_loss)
+        return {'loss': combined_loss}
+
+    def validation_step(self, batch, batch_idx):
+        x, y = batch
+        y_hat = torch.sigmoid(self(x))
+
+        wce_loss = self.wce(y_hat, y)
+        dice = dice_loss(y_hat, y)
+
+        combined_loss = (1 - self.alpha) * wce_loss + self.alpha * dice
+
+        precision = self.val_precision(y_hat, y)
+        recall = self.val_recall(y_hat, y)
+        iou = self.val_iou(y_hat, y)
+        loss_f1 = self.val_f1(y_hat, y)
+
+        self.log('val_precision', precision)
+        self.log('val_recall', recall)
+        self.log('val_wce', wce_loss)
+        self.log('val_dice', dice)
+        self.log('val_iou', iou)
+        self.log('val_f1', loss_f1)
+        self.log('val_loss', combined_loss)
+
+        if self.current_epoch % 10 == 0:
+            x = (x - x.min()) / (x.max() - x.min())
+            x = x[:, 0:3]
+            x = x.permute(0, 2, 3, 1)
+            y_hat = (y_hat > 0.5).float()
+
+            class_labels = {0: "no landslide", 1: "landslide"}  # Define class_labels here
+
+            self.logger.experiment.log({
+                "image": wandb.Image(x[0].cpu().detach().numpy(), masks={ 
+                    "predictions": {
+                        "mask_data": y_hat[0][0].cpu().detach().numpy(),
+                        "class_labels": class_labels
+                    },
+                    "ground_truth": {
+                        "mask_data": y[0][0].cpu().detach().numpy(),
+                        "class_labels": class_labels
+                    }
+                })
+            })
+        return {'val_loss': combined_loss}
+
+    def configure_optimizers(self):
+        optimizer = Adam(self.parameters(), lr=self.lr)
+        scheduler = StepLR(optimizer, step_size=30, gamma=0.1)
+        return [optimizer], [scheduler]
+
+class Block(nn.Module):
+    def __init__(self, inputs=3, middles=64, outs=64):
+        super().__init__()
+
+        self.conv1 = nn.Conv2d(inputs, middles, 3, 1, 1)
+        self.conv2 = nn.Conv2d(middles, outs, 3, 1, 1)
+        self.relu = nn.ReLU()
+        self.bn = nn.BatchNorm2d(outs)
+        self.pool = nn.MaxPool2d(2, 2)
+
+    def forward(self, x):
+        x = self.relu(self.conv1(x))
+        x = self.relu(self.bn(self.conv2(x)))
+        return self.pool(x), x
+
+class UNet(nn.Module):
+    def __init__(self, in_channels=3, out_channels=1):
+        super().__init__()
+
+        self.en1 = Block(in_channels, 64, 64)
+        self.en2 = Block(64, 128, 128)
+        self.en3 = Block(128, 256, 256)
+        self.en4 = Block(256, 512, 512)
+        self.en5 = Block(512, 1024, 512)
+
+        self.upsample4 = nn.ConvTranspose2d(512, 512, 2, stride=2)
+        self.de4 = Block(1024, 512, 256)
+
+        self.upsample3 = nn.ConvTranspose2d(256, 256, 2, stride=2)
+        self.de3 = Block(512, 256, 128)
+
+        self.upsample2 = nn.ConvTranspose2d(128, 128, 2, stride=2)
+        self.de2 = Block(256, 128, 64)
+
+        self.upsample1 = nn.ConvTranspose2d(64, 64, 2, stride=2)
+        self.de1 = Block(128, 64, 64)
+
+        self.conv_last = nn.Conv2d(64, out_channels, kernel_size=1, stride=1, padding=0)
+
+    def forward(self, x):
+        x, e1 = self.en1(x)
+        x, e2 = self.en2(x)
+        x, e3 = self.en3(x)
+        x, e4 = self.en4(x)
+        _, x = self.en5(x)
+
+        x = self.upsample4(x)
+        x = torch.cat([x, e4], dim=1)
+        _, x = self.de4(x)
+
+        x = self.upsample3(x)
+        x = torch.cat([x, e3], dim=1)
+        _, x = self.de3(x)
+
+        x = self.upsample2(x)
+        x = torch.cat([x, e2], dim=1)
+        _, x = self.de2(x)
+
+        x = self.upsample1(x)
+        x = torch.cat([x, e1], dim=1)
+        _, x = self.de1(x)
+
+        x = self.conv_last(x)
+
+        return x
+
+def dice_loss(y_hat, y):
+    smooth = 1e-6
+    y_hat = y_hat.view(-1)
+    y = y.view(-1)
+    intersection = (y_hat * y).sum()
+    union = y_hat.sum() + y.sum()
+    dice = (2 * intersection + smooth) / (union + smooth)
+    return 1 - dice

src/resnet34_model.py

@@ -0,0 +1,209 @@
+import torch
+import torch.nn as nn
+import segmentation_models_pytorch as smp
+from torchmetrics import F1Score, Precision, Recall, JaccardIndex
+import pytorch_lightning as pl
+import wandb
+from torch.optim import Adam
+from torch.optim.lr_scheduler import StepLR
+
+class smp_model_resnet34(nn.Module):
+    def __init__(self, in_channels, out_channels, model_type, num_classes, encoder_weights):
+        super(smp_model_resnet34, self).__init__()
+        
+        # Use ResNet34 pre-trained model as encoder
+        self.model = smp.Unet(
+            encoder_name='resnet34',
+            encoder_weights=encoder_weights,
+            in_channels=in_channels,  # The number of input channels, which is 14
+            classes=num_classes,       # Output classes, which is 1
+        )
+
+    def forward(self, x):
+        return self.model(x)
+
+class LandslideModel(pl.LightningModule):
+    def __init__(self, config, alpha=0.5):
+        super(LandslideModel, self).__init__()
+
+        model_type = config['model_config']['model_type']
+        in_channels = config['model_config']['in_channels']
+        num_classes = config['model_config']['num_classes']
+        self.alpha = alpha  # Assign the alpha value to the class variable
+        self.lr = config['train_config']['lr']
+
+        if model_type == 'unet':
+            self.model = UNet(in_channels=in_channels, out_channels=num_classes)
+        else:
+            encoder_weights = config['model_config']['encoder_weights']
+            # Use the custom smp_model_resnet34 instead of smp_model
+            self.model = smp_model_resnet34(in_channels=in_channels, 
+                                            out_channels=num_classes, 
+                                            model_type=model_type, 
+                                            num_classes=num_classes, 
+                                            encoder_weights=encoder_weights)
+
+        self.weights = torch.tensor([5], dtype=torch.float32).to(self.device)
+        self.wce = nn.BCELoss(weight=self.weights)
+
+        self.train_f1 = F1Score(task='binary')
+        self.val_f1 = F1Score(task='binary')
+
+        self.train_precision = Precision(task='binary')
+        self.val_precision = Precision(task='binary')
+
+        self.train_recall = Recall(task='binary')
+        self.val_recall = Recall(task='binary')
+
+        self.train_iou = JaccardIndex(task='binary')
+        self.val_iou = JaccardIndex(task='binary')
+
+    def forward(self, x):
+        return self.model(x)
+
+    def training_step(self, batch, batch_idx):
+        x, y = batch
+        y_hat = torch.sigmoid(self(x))
+
+        wce_loss = self.wce(y_hat, y)
+        dice = dice_loss(y_hat, y)
+
+        combined_loss = (1 - self.alpha) * wce_loss + self.alpha * dice
+
+        precision = self.train_precision(y_hat, y)
+        recall = self.train_recall(y_hat, y)
+        iou = self.train_iou(y_hat, y)
+        loss_f1 = self.train_f1(y_hat, y)
+
+        self.log('train_precision', precision)
+        self.log('train_recall', recall)
+        self.log('train_wce', wce_loss)
+        self.log('train_dice', dice)
+        self.log('train_iou', iou)
+        self.log('train_f1', loss_f1)
+        self.log('train_loss', combined_loss)
+        return {'loss': combined_loss}
+
+    def validation_step(self, batch, batch_idx):
+        x, y = batch
+        y_hat = torch.sigmoid(self(x))
+
+        wce_loss = self.wce(y_hat, y)
+        dice = dice_loss(y_hat, y)
+
+        combined_loss = (1 - self.alpha) * wce_loss + self.alpha * dice
+
+        precision = self.val_precision(y_hat, y)
+        recall = self.val_recall(y_hat, y)
+        iou = self.val_iou(y_hat, y)
+        loss_f1 = self.val_f1(y_hat, y)
+
+        self.log('val_precision', precision)
+        self.log('val_recall', recall)
+        self.log('val_wce', wce_loss)
+        self.log('val_dice', dice)
+        self.log('val_iou', iou)
+        self.log('val_f1', loss_f1)
+        self.log('val_loss', combined_loss)
+
+        if self.current_epoch % 10 == 0:
+            x = (x - x.min()) / (x.max() - x.min())
+            x = x[:, 0:3]
+            x = x.permute(0, 2, 3, 1)
+            y_hat = (y_hat > 0.5).float()
+
+            class_labels = {0: "no landslide", 1: "landslide"}  # Define class_labels here
+
+            self.logger.experiment.log({
+                "image": wandb.Image(x[0].cpu().detach().numpy(), masks={ 
+                    "predictions": {
+                        "mask_data": y_hat[0][0].cpu().detach().numpy(),
+                        "class_labels": class_labels
+                    },
+                    "ground_truth": {
+                        "mask_data": y[0][0].cpu().detach().numpy(),
+                        "class_labels": class_labels
+                    }
+                })
+            })
+        return {'val_loss': combined_loss}
+
+    def configure_optimizers(self):
+        optimizer = Adam(self.parameters(), lr=self.lr)
+        scheduler = StepLR(optimizer, step_size=30, gamma=0.1)
+        return [optimizer], [scheduler]
+
+class Block(nn.Module):
+    def __init__(self, inputs=3, middles=64, outs=64):
+        super().__init__()
+
+        self.conv1 = nn.Conv2d(inputs, middles, 3, 1, 1)
+        self.conv2 = nn.Conv2d(middles, outs, 3, 1, 1)
+        self.relu = nn.ReLU()
+        self.bn = nn.BatchNorm2d(outs)
+        self.pool = nn.MaxPool2d(2, 2)
+
+    def forward(self, x):
+        x = self.relu(self.conv1(x))
+        x = self.relu(self.bn(self.conv2(x)))
+        return self.pool(x), x
+
+class UNet(nn.Module):
+    def __init__(self, in_channels=3, out_channels=1):
+        super().__init__()
+
+        self.en1 = Block(in_channels, 64, 64)
+        self.en2 = Block(64, 128, 128)
+        self.en3 = Block(128, 256, 256)
+        self.en4 = Block(256, 512, 512)
+        self.en5 = Block(512, 1024, 512)
+
+        self.upsample4 = nn.ConvTranspose2d(512, 512, 2, stride=2)
+        self.de4 = Block(1024, 512, 256)
+
+        self.upsample3 = nn.ConvTranspose2d(256, 256, 2, stride=2)
+        self.de3 = Block(512, 256, 128)
+
+        self.upsample2 = nn.ConvTranspose2d(128, 128, 2, stride=2)
+        self.de2 = Block(256, 128, 64)
+
+        self.upsample1 = nn.ConvTranspose2d(64, 64, 2, stride=2)
+        self.de1 = Block(128, 64, 64)
+
+        self.conv_last = nn.Conv2d(64, out_channels, kernel_size=1, stride=1, padding=0)
+
+    def forward(self, x):
+        x, e1 = self.en1(x)
+        x, e2 = self.en2(x)
+        x, e3 = self.en3(x)
+        x, e4 = self.en4(x)
+        _, x = self.en5(x)
+
+        x = self.upsample4(x)
+        x = torch.cat([x, e4], dim=1)
+        _, x = self.de4(x)
+
+        x = self.upsample3(x)
+        x = torch.cat([x, e3], dim=1)
+        _, x = self.de3(x)
+
+        x = self.upsample2(x)
+        x = torch.cat([x, e2], dim=1)
+        _, x = self.de2(x)
+
+        x = self.upsample1(x)
+        x = torch.cat([x, e1], dim=1)
+        _, x = self.de1(x)
+
+        x = self.conv_last(x)
+
+        return x
+
+def dice_loss(y_hat, y):
+    smooth = 1e-6
+    y_hat = y_hat.view(-1)
+    y = y.view(-1)
+    intersection = (y_hat * y).sum()
+    union = y_hat.sum() + y.sum()
+    dice = (2 * intersection + smooth) / (union + smooth)
+    return 1 - dice

src/resnext50_32x4d_model.py

@@ -0,0 +1,207 @@
+import torch
+import torch.nn as nn
+import segmentation_models_pytorch as smp
+from torchmetrics import F1Score, Precision, Recall, JaccardIndex
+import pytorch_lightning as pl
+import wandb
+from torch.optim import Adam
+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__()
+        self.model = smp.Unet(
+            encoder_name=model_type,
+            encoder_weights=encoder_weights,
+            in_channels=in_channels,  # Use the original in_channels
+            classes=num_classes,
+        )
+
+    def forward(self, x):
+        x = self.model(x)
+        return x
+
+class LandslideModel(pl.LightningModule):
+    def __init__(self, config, alpha=0.5):
+        super(LandslideModel, self).__init__()
+
+        model_type = config['model_config']['model_type']
+        in_channels = config['model_config']['in_channels']
+        num_classes = config['model_config']['num_classes']
+        self.alpha = alpha
+        self.lr = config['train_config']['lr']
+
+        if model_type == 'unet':
+            self.model = UNet(in_channels=in_channels, out_channels=num_classes)
+        else:
+            encoder_weights = config['model_config']['encoder_weights']
+            self.model = smp_model(in_channels=in_channels, 
+                                   out_channels=num_classes, 
+                                   model_type=model_type, 
+                                   num_classes=num_classes, 
+                                   encoder_weights=encoder_weights)
+
+        self.weights = torch.tensor([5], dtype=torch.float32).to(self.device)
+        self.wce = nn.BCELoss(weight=self.weights)
+
+        self.train_f1 = F1Score(task='binary')
+        self.val_f1 = F1Score(task='binary')
+
+        self.train_precision = Precision(task='binary')
+        self.val_precision = Precision(task='binary')
+
+        self.train_recall = Recall(task='binary')
+        self.val_recall = Recall(task='binary')
+
+        self.train_iou = JaccardIndex(task='binary')
+        self.val_iou = JaccardIndex(task='binary')
+
+    def forward(self, x):
+        return self.model(x)
+
+    def training_step(self, batch, batch_idx):
+        x, y = batch
+        y_hat = torch.sigmoid(self(x))
+
+        wce_loss = self.wce(y_hat, y)
+        dice = dice_loss(y_hat, y)
+
+        combined_loss = (1 - self.alpha) * wce_loss + self.alpha * dice
+
+        precision = self.train_precision(y_hat, y)
+        recall = self.train_recall(y_hat, y)
+        iou = self.train_iou(y_hat, y)
+        loss_f1 = self.train_f1(y_hat, y)
+
+        self.log('train_precision', precision)
+        self.log('train_recall', recall)
+        self.log('train_wce', wce_loss)
+        self.log('train_dice', dice)
+        self.log('train_iou', iou)
+        self.log('train_f1', loss_f1)
+        self.log('train_loss', combined_loss)
+        return {'loss': combined_loss}
+
+    def validation_step(self, batch, batch_idx):
+        x, y = batch
+        y_hat = torch.sigmoid(self(x))
+
+        wce_loss = self.wce(y_hat, y)
+        dice = dice_loss(y_hat, y)
+
+        combined_loss = (1 - self.alpha) * wce_loss + self.alpha * dice
+
+        precision = self.val_precision(y_hat, y)
+        recall = self.val_recall(y_hat, y)
+        iou = self.val_iou(y_hat, y)
+        loss_f1 = self.val_f1(y_hat, y)
+
+        self.log('val_precision', precision)
+        self.log('val_recall', recall)
+        self.log('val_wce', wce_loss)
+        self.log('val_dice', dice)
+        self.log('val_iou', iou)
+        self.log('val_f1', loss_f1)
+        self.log('val_loss', combined_loss)
+
+        if self.current_epoch % 10 == 0:
+            x = (x - x.min()) / (x.max() - x.min())
+            x = x[:, 0:3]
+            x = x.permute(0, 2, 3, 1)
+            y_hat = (y_hat > 0.5).float()
+
+            class_labels = {0: "no landslide", 1: "landslide"}
+
+            self.logger.experiment.log({
+                "image": wandb.Image(x[0].cpu().detach().numpy(), masks={ 
+                    "predictions": {
+                        "mask_data": y_hat[0][0].cpu().detach().numpy(),
+                        "class_labels": class_labels
+                    },
+                    "ground_truth": {
+                        "mask_data": y[0][0].cpu().detach().numpy(),
+                        "class_labels": class_labels
+                    }
+                })
+            })
+        return {'val_loss': combined_loss}
+
+    def configure_optimizers(self):
+        optimizer = Adam(self.parameters(), lr=self.lr)
+        scheduler = StepLR(optimizer, step_size=30, gamma=0.1)
+        return [optimizer], [scheduler]
+
+class Block(nn.Module):
+    def __init__(self, inputs=3, middles=64, outs=64):
+        super().__init__()
+
+        self.conv1 = nn.Conv2d(inputs, middles, 3, 1, 1)
+        self.conv2 = nn.Conv2d(middles, outs, 3, 1, 1)
+        self.relu = nn.ReLU()
+        self.bn = nn.BatchNorm2d(outs)
+        self.pool = nn.MaxPool2d(2, 2)
+
+    def forward(self, x):
+        x = self.relu(self.conv1(x))
+        x = self.relu(self.bn(self.conv2(x)))
+        return self.pool(x), x
+
+class UNet(nn.Module):
+    def __init__(self, in_channels=3, out_channels=1):
+        super().__init__()
+
+        self.en1 = Block(in_channels, 64, 64)
+        self.en2 = Block(64, 128, 128)
+        self.en3 = Block(128, 256, 256)
+        self.en4 = Block(256, 512, 512)
+        self.en5 = Block(512, 1024, 512)
+
+        self.upsample4 = nn.ConvTranspose2d(512, 512, 2, stride=2)
+        self.de4 = Block(1024, 512, 256)
+
+        self.upsample3 = nn.ConvTranspose2d(256, 256, 2, stride=2)
+        self.de3 = Block(512, 256, 128)
+
+        self.upsample2 = nn.ConvTranspose2d(128, 128, 2, stride=2)
+        self.de2 = Block(256, 128, 64)
+
+        self.upsample1 = nn.ConvTranspose2d(64, 64, 2, stride=2)
+        self.de1 = Block(128, 64, 64)
+
+        self.conv_last = nn.Conv2d(64, out_channels, kernel_size=1, stride=1, padding=0)
+
+    def forward(self, x):
+        x, e1 = self.en1(x)
+        x, e2 = self.en2(x)
+        x, e3 = self.en3(x)
+        x, e4 = self.en4(x)
+        _, x = self.en5(x)
+
+        x = self.upsample4(x)
+        x = torch.cat([x, e4], dim=1)
+        _, x = self.de4(x)
+
+        x = self.upsample3(x)
+        x = torch.cat([x, e3], dim=1)
+        _, x = self.de3(x)
+
+        x = self.upsample2(x)
+        x = torch.cat([x, e2], dim=1)
+        _, x = self.de2(x)
+
+        x = self.upsample1(x)
+        x = torch.cat([x, e1], dim=1)
+        _, x = self.de1(x)
+
+        x = self.conv_last(x)
+
+        return x
+
+def dice_loss(y_hat, y):
+    smooth = 1e-6
+    y_hat = y_hat.view(-1)
+    y = y.view(-1)
+    intersection = (y_hat * y).sum()
+    union = y_hat.sum() + y.sum()
+    dice = (2 * intersection + smooth) / (union + smooth)
+    return 1 - dice

src/se_resnet50_model.py

@@ -0,0 +1,207 @@
+import torch
+import torch.nn as nn
+import segmentation_models_pytorch as smp
+from torchmetrics import F1Score, Precision, Recall, JaccardIndex
+import pytorch_lightning as pl
+import wandb
+from torch.optim import Adam
+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__()
+        self.model = smp.Unet(
+            encoder_name=model_type,
+            encoder_weights=encoder_weights,
+            in_channels=in_channels,  # Use the original in_channels
+            classes=num_classes,
+        )
+
+    def forward(self, x):
+        x = self.model(x)
+        return x
+
+class LandslideModel(pl.LightningModule):
+    def __init__(self, config, alpha=0.5):
+        super(LandslideModel, self).__init__()
+
+        model_type = config['model_config']['model_type']
+        in_channels = config['model_config']['in_channels']
+        num_classes = config['model_config']['num_classes']
+        self.alpha = alpha
+        self.lr = config['train_config']['lr']
+
+        if model_type == 'unet':
+            self.model = UNet(in_channels=in_channels, out_channels=num_classes)
+        else:
+            encoder_weights = config['model_config']['encoder_weights']
+            self.model = smp_model(in_channels=in_channels, 
+                                   out_channels=num_classes, 
+                                   model_type=model_type, 
+                                   num_classes=num_classes, 
+                                   encoder_weights=encoder_weights)
+
+        self.weights = torch.tensor([5], dtype=torch.float32).to(self.device)
+        self.wce = nn.BCELoss(weight=self.weights)
+
+        self.train_f1 = F1Score(task='binary')
+        self.val_f1 = F1Score(task='binary')
+
+        self.train_precision = Precision(task='binary')
+        self.val_precision = Precision(task='binary')
+
+        self.train_recall = Recall(task='binary')
+        self.val_recall = Recall(task='binary')
+
+        self.train_iou = JaccardIndex(task='binary')
+        self.val_iou = JaccardIndex(task='binary')
+
+    def forward(self, x):
+        return self.model(x)
+
+    def training_step(self, batch, batch_idx):
+        x, y = batch
+        y_hat = torch.sigmoid(self(x))
+
+        wce_loss = self.wce(y_hat, y)
+        dice = dice_loss(y_hat, y)
+
+        combined_loss = (1 - self.alpha) * wce_loss + self.alpha * dice
+
+        precision = self.train_precision(y_hat, y)
+        recall = self.train_recall(y_hat, y)
+        iou = self.train_iou(y_hat, y)
+        loss_f1 = self.train_f1(y_hat, y)
+
+        self.log('train_precision', precision)
+        self.log('train_recall', recall)
+        self.log('train_wce', wce_loss)
+        self.log('train_dice', dice)
+        self.log('train_iou', iou)
+        self.log('train_f1', loss_f1)
+        self.log('train_loss', combined_loss)
+        return {'loss': combined_loss}
+
+    def validation_step(self, batch, batch_idx):
+        x, y = batch
+        y_hat = torch.sigmoid(self(x))
+
+        wce_loss = self.wce(y_hat, y)
+        dice = dice_loss(y_hat, y)
+
+        combined_loss = (1 - self.alpha) * wce_loss + self.alpha * dice
+
+        precision = self.val_precision(y_hat, y)
+        recall = self.val_recall(y_hat, y)
+        iou = self.val_iou(y_hat, y)
+        loss_f1 = self.val_f1(y_hat, y)
+
+        self.log('val_precision', precision)
+        self.log('val_recall', recall)
+        self.log('val_wce', wce_loss)
+        self.log('val_dice', dice)
+        self.log('val_iou', iou)
+        self.log('val_f1', loss_f1)
+        self.log('val_loss', combined_loss)
+
+        if self.current_epoch % 10 == 0:
+            x = (x - x.min()) / (x.max() - x.min())
+            x = x[:, 0:3]
+            x = x.permute(0, 2, 3, 1)
+            y_hat = (y_hat > 0.5).float()
+
+            class_labels = {0: "no landslide", 1: "landslide"}
+
+            self.logger.experiment.log({
+                "image": wandb.Image(x[0].cpu().detach().numpy(), masks={ 
+                    "predictions": {
+                        "mask_data": y_hat[0][0].cpu().detach().numpy(),
+                        "class_labels": class_labels
+                    },
+                    "ground_truth": {
+                        "mask_data": y[0][0].cpu().detach().numpy(),
+                        "class_labels": class_labels
+                    }
+                })
+            })
+        return {'val_loss': combined_loss}
+
+    def configure_optimizers(self):
+        optimizer = Adam(self.parameters(), lr=self.lr)
+        scheduler = StepLR(optimizer, step_size=30, gamma=0.1)
+        return [optimizer], [scheduler]
+
+class Block(nn.Module):
+    def __init__(self, inputs=3, middles=64, outs=64):
+        super().__init__()
+
+        self.conv1 = nn.Conv2d(inputs, middles, 3, 1, 1)
+        self.conv2 = nn.Conv2d(middles, outs, 3, 1, 1)
+        self.relu = nn.ReLU()
+        self.bn = nn.BatchNorm2d(outs)
+        self.pool = nn.MaxPool2d(2, 2)
+
+    def forward(self, x):
+        x = self.relu(self.conv1(x))
+        x = self.relu(self.bn(self.conv2(x)))
+        return self.pool(x), x
+
+class UNet(nn.Module):
+    def __init__(self, in_channels=3, out_channels=1):
+        super().__init__()
+
+        self.en1 = Block(in_channels, 64, 64)
+        self.en2 = Block(64, 128, 128)
+        self.en3 = Block(128, 256, 256)
+        self.en4 = Block(256, 512, 512)
+        self.en5 = Block(512, 1024, 512)
+
+        self.upsample4 = nn.ConvTranspose2d(512, 512, 2, stride=2)
+        self.de4 = Block(1024, 512, 256)
+
+        self.upsample3 = nn.ConvTranspose2d(256, 256, 2, stride=2)
+        self.de3 = Block(512, 256, 128)
+
+        self.upsample2 = nn.ConvTranspose2d(128, 128, 2, stride=2)
+        self.de2 = Block(256, 128, 64)
+
+        self.upsample1 = nn.ConvTranspose2d(64, 64, 2, stride=2)
+        self.de1 = Block(128, 64, 64)
+
+        self.conv_last = nn.Conv2d(64, out_channels, kernel_size=1, stride=1, padding=0)
+
+    def forward(self, x):
+        x, e1 = self.en1(x)
+        x, e2 = self.en2(x)
+        x, e3 = self.en3(x)
+        x, e4 = self.en4(x)
+        _, x = self.en5(x)
+
+        x = self.upsample4(x)
+        x = torch.cat([x, e4], dim=1)
+        _, x = self.de4(x)
+
+        x = self.upsample3(x)
+        x = torch.cat([x, e3], dim=1)
+        _, x = self.de3(x)
+
+        x = self.upsample2(x)
+        x = torch.cat([x, e2], dim=1)
+        _, x = self.de2(x)
+
+        x = self.upsample1(x)
+        x = torch.cat([x, e1], dim=1)
+        _, x = self.de1(x)
+
+        x = self.conv_last(x)
+
+        return x
+
+def dice_loss(y_hat, y):
+    smooth = 1e-6
+    y_hat = y_hat.view(-1)
+    y = y.view(-1)
+    intersection = (y_hat * y).sum()
+    union = y_hat.sum() + y.sum()
+    dice = (2 * intersection + smooth) / (union + smooth)
+    return 1 - dice

src/se_resnext50_32x4d_model.py

@@ -0,0 +1,207 @@
+import torch
+import torch.nn as nn
+import segmentation_models_pytorch as smp
+from torchmetrics import F1Score, Precision, Recall, JaccardIndex
+import pytorch_lightning as pl
+import wandb
+from torch.optim import Adam
+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__()
+        self.model = smp.Unet(
+            encoder_name=model_type,
+            encoder_weights=encoder_weights,
+            in_channels=in_channels,  # Use the original in_channels
+            classes=num_classes,
+        )
+
+    def forward(self, x):
+        x = self.model(x)
+        return x
+
+class LandslideModel(pl.LightningModule):
+    def __init__(self, config, alpha=0.5):
+        super(LandslideModel, self).__init__()
+
+        model_type = config['model_config']['model_type']
+        in_channels = config['model_config']['in_channels']
+        num_classes = config['model_config']['num_classes']
+        self.alpha = alpha
+        self.lr = config['train_config']['lr']
+
+        if model_type == 'unet':
+            self.model = UNet(in_channels=in_channels, out_channels=num_classes)
+        else:
+            encoder_weights = config['model_config']['encoder_weights']
+            self.model = smp_model(in_channels=in_channels, 
+                                   out_channels=num_classes, 
+                                   model_type=model_type, 
+                                   num_classes=num_classes, 
+                                   encoder_weights=encoder_weights)
+
+        self.weights = torch.tensor([5], dtype=torch.float32).to(self.device)
+        self.wce = nn.BCELoss(weight=self.weights)
+
+        self.train_f1 = F1Score(task='binary')
+        self.val_f1 = F1Score(task='binary')
+
+        self.train_precision = Precision(task='binary')
+        self.val_precision = Precision(task='binary')
+
+        self.train_recall = Recall(task='binary')
+        self.val_recall = Recall(task='binary')
+
+        self.train_iou = JaccardIndex(task='binary')
+        self.val_iou = JaccardIndex(task='binary')
+
+    def forward(self, x):
+        return self.model(x)
+
+    def training_step(self, batch, batch_idx):
+        x, y = batch
+        y_hat = torch.sigmoid(self(x))
+
+        wce_loss = self.wce(y_hat, y)
+        dice = dice_loss(y_hat, y)
+
+        combined_loss = (1 - self.alpha) * wce_loss + self.alpha * dice
+
+        precision = self.train_precision(y_hat, y)
+        recall = self.train_recall(y_hat, y)
+        iou = self.train_iou(y_hat, y)
+        loss_f1 = self.train_f1(y_hat, y)
+
+        self.log('train_precision', precision)
+        self.log('train_recall', recall)
+        self.log('train_wce', wce_loss)
+        self.log('train_dice', dice)
+        self.log('train_iou', iou)
+        self.log('train_f1', loss_f1)
+        self.log('train_loss', combined_loss)
+        return {'loss': combined_loss}
+
+    def validation_step(self, batch, batch_idx):
+        x, y = batch
+        y_hat = torch.sigmoid(self(x))
+
+        wce_loss = self.wce(y_hat, y)
+        dice = dice_loss(y_hat, y)
+
+        combined_loss = (1 - self.alpha) * wce_loss + self.alpha * dice
+
+        precision = self.val_precision(y_hat, y)
+        recall = self.val_recall(y_hat, y)
+        iou = self.val_iou(y_hat, y)
+        loss_f1 = self.val_f1(y_hat, y)
+
+        self.log('val_precision', precision)
+        self.log('val_recall', recall)
+        self.log('val_wce', wce_loss)
+        self.log('val_dice', dice)
+        self.log('val_iou', iou)
+        self.log('val_f1', loss_f1)
+        self.log('val_loss', combined_loss)
+
+        if self.current_epoch % 10 == 0:
+            x = (x - x.min()) / (x.max() - x.min())
+            x = x[:, 0:3]
+            x = x.permute(0, 2, 3, 1)
+            y_hat = (y_hat > 0.5).float()
+
+            class_labels = {0: "no landslide", 1: "landslide"}
+
+            self.logger.experiment.log({
+                "image": wandb.Image(x[0].cpu().detach().numpy(), masks={ 
+                    "predictions": {
+                        "mask_data": y_hat[0][0].cpu().detach().numpy(),
+                        "class_labels": class_labels
+                    },
+                    "ground_truth": {
+                        "mask_data": y[0][0].cpu().detach().numpy(),
+                        "class_labels": class_labels
+                    }
+                })
+            })
+        return {'val_loss': combined_loss}
+
+    def configure_optimizers(self):
+        optimizer = Adam(self.parameters(), lr=self.lr)
+        scheduler = StepLR(optimizer, step_size=30, gamma=0.1)
+        return [optimizer], [scheduler]
+
+class Block(nn.Module):
+    def __init__(self, inputs=3, middles=64, outs=64):
+        super().__init__()
+
+        self.conv1 = nn.Conv2d(inputs, middles, 3, 1, 1)
+        self.conv2 = nn.Conv2d(middles, outs, 3, 1, 1)
+        self.relu = nn.ReLU()
+        self.bn = nn.BatchNorm2d(outs)
+        self.pool = nn.MaxPool2d(2, 2)
+
+    def forward(self, x):
+        x = self.relu(self.conv1(x))
+        x = self.relu(self.bn(self.conv2(x)))
+        return self.pool(x), x
+
+class UNet(nn.Module):
+    def __init__(self, in_channels=3, out_channels=1):
+        super().__init__()
+
+        self.en1 = Block(in_channels, 64, 64)
+        self.en2 = Block(64, 128, 128)
+        self.en3 = Block(128, 256, 256)
+        self.en4 = Block(256, 512, 512)
+        self.en5 = Block(512, 1024, 512)
+
+        self.upsample4 = nn.ConvTranspose2d(512, 512, 2, stride=2)
+        self.de4 = Block(1024, 512, 256)
+
+        self.upsample3 = nn.ConvTranspose2d(256, 256, 2, stride=2)
+        self.de3 = Block(512, 256, 128)
+
+        self.upsample2 = nn.ConvTranspose2d(128, 128, 2, stride=2)
+        self.de2 = Block(256, 128, 64)
+
+        self.upsample1 = nn.ConvTranspose2d(64, 64, 2, stride=2)
+        self.de1 = Block(128, 64, 64)
+
+        self.conv_last = nn.Conv2d(64, out_channels, kernel_size=1, stride=1, padding=0)
+
+    def forward(self, x):
+        x, e1 = self.en1(x)
+        x, e2 = self.en2(x)
+        x, e3 = self.en3(x)
+        x, e4 = self.en4(x)
+        _, x = self.en5(x)
+
+        x = self.upsample4(x)
+        x = torch.cat([x, e4], dim=1)
+        _, x = self.de4(x)
+
+        x = self.upsample3(x)
+        x = torch.cat([x, e3], dim=1)
+        _, x = self.de3(x)
+
+        x = self.upsample2(x)
+        x = torch.cat([x, e2], dim=1)
+        _, x = self.de2(x)
+
+        x = self.upsample1(x)
+        x = torch.cat([x, e1], dim=1)
+        _, x = self.de1(x)
+
+        x = self.conv_last(x)
+
+        return x
+
+def dice_loss(y_hat, y):
+    smooth = 1e-6
+    y_hat = y_hat.view(-1)
+    y = y.view(-1)
+    intersection = (y_hat * y).sum()
+    union = y_hat.sum() + y.sum()
+    dice = (2 * intersection + smooth) / (union + smooth)
+    return 1 - dice

src/segformer_model.py

@@ -0,0 +1,138 @@
+import torch
+import torch.nn as nn
+import torchmetrics
+import pytorch_lightning as pl
+import wandb
+from torch.optim import Adam
+from torch.optim.lr_scheduler import StepLR
+from transformers import SegformerForSemanticSegmentation
+
+class LandslideModel(pl.LightningModule):
+    def __init__(self, config, alpha=0.5):
+        super(LandslideModel, self).__init__()
+
+        self.model_type = config['model_config']['model_type']
+        self.in_channels = config['model_config']['in_channels']
+        self.num_classes = config['model_config']['num_classes']
+        self.alpha = alpha
+        self.lr = config['train_config']['lr']
+
+        self.model = SegformerForSemanticSegmentation.from_pretrained(
+            "nvidia/segformer-b2-finetuned-ade-512-512",
+            ignore_mismatched_sizes=True,
+            num_labels=self.num_classes
+        )
+
+        # Modify the input layer for 14 channels
+        self.model.segformer.encoder.patch_embeddings[0].proj = nn.Conv2d(
+            in_channels=self.in_channels,
+            out_channels=self.model.segformer.encoder.patch_embeddings[0].proj.out_channels,
+            kernel_size=self.model.segformer.encoder.patch_embeddings[0].proj.kernel_size,
+            stride=self.model.segformer.encoder.patch_embeddings[0].proj.stride,
+            padding=self.model.segformer.encoder.patch_embeddings[0].proj.padding
+        )
+
+        self.weights = torch.tensor([5], dtype=torch.float32).to(self.device)
+        self.wce = nn.BCELoss(weight=self.weights)
+
+        self.train_f1 = torchmetrics.F1Score(task='binary')
+        self.val_f1 = torchmetrics.F1Score(task='binary')
+
+        self.train_precision = torchmetrics.Precision(task='binary')
+        self.val_precision = torchmetrics.Precision(task='binary')
+
+        self.train_recall = torchmetrics.Recall(task='binary')
+        self.val_recall = torchmetrics.Recall(task='binary')
+
+        self.train_iou = torchmetrics.JaccardIndex(task='binary')
+        self.val_iou = torchmetrics.JaccardIndex(task='binary')
+
+    def forward(self, x):
+        return self.model(x).logits
+
+    def training_step(self, batch, batch_idx):
+        x, y = batch
+        y_hat = torch.sigmoid(self(x))
+
+        # Resize y_hat to match the size of y
+        y_hat = nn.functional.interpolate(y_hat, size=y.shape[2:], mode='bilinear', align_corners=False)
+
+        wce_loss = self.wce(y_hat, y)
+        dice = dice_loss(y_hat, y)
+
+        combined_loss = (1 - self.alpha) * wce_loss + self.alpha * dice
+
+        precision = self.train_precision(y_hat, y)
+        recall = self.train_recall(y_hat, y)
+        iou = self.train_iou(y_hat, y)
+        loss_f1 = self.train_f1(y_hat, y)
+
+        self.log('train_precision', precision)
+        self.log('train_recall', recall)
+        self.log('train_wce', wce_loss)
+        self.log('train_dice', dice)
+        self.log('train_iou', iou)
+        self.log('train_f1', loss_f1)
+        self.log('train_loss', combined_loss)
+        return {'loss': combined_loss}
+
+    def validation_step(self, batch, batch_idx):
+        x, y = batch
+        y_hat = torch.sigmoid(self(x))
+
+        # Resize y_hat to match the size of y
+        y_hat = nn.functional.interpolate(y_hat, size=y.shape[2:], mode='bilinear', align_corners=False)
+
+        wce_loss = self.wce(y_hat, y)
+        dice = dice_loss(y_hat, y)
+
+        combined_loss = (1 - self.alpha) * wce_loss + self.alpha * dice
+
+        precision = self.val_precision(y_hat, y)
+        recall = self.val_recall(y_hat, y)
+        iou = self.val_iou(y_hat, y)
+        loss_f1 = self.val_f1(y_hat, y)
+
+        self.log('val_precision', precision)
+        self.log('val_recall', recall)
+        self.log('val_wce', wce_loss)
+        self.log('val_dice', dice)
+        self.log('val_iou', iou)
+        self.log('val_f1', loss_f1)
+        self.log('val_loss', combined_loss)
+
+        if self.current_epoch % 10 == 0:
+            x = (x - x.min()) / (x.max() - x.min())
+            x = x[:, 0:3]
+            x = x.permute(0, 2, 3, 1)
+            y_hat = (y_hat > 0.5).float()
+
+            class_labels = {0: "no landslide", 1: "landslide"}
+
+            self.logger.experiment.log({
+                "image": wandb.Image(x[0].cpu().detach().numpy(), masks={ 
+                    "predictions": {
+                        "mask_data": y_hat[0][0].cpu().detach().numpy(),
+                        "class_labels": class_labels
+                    },
+                    "ground_truth": {
+                        "mask_data": y[0][0].cpu().detach().numpy(),
+                        "class_labels": class_labels
+                    }
+                })
+            })
+        return {'val_loss': combined_loss}
+
+    def configure_optimizers(self):
+        optimizer = Adam(self.parameters(), lr=self.lr)
+        scheduler = StepLR(optimizer, step_size=30, gamma=0.1)
+        return [optimizer], [scheduler]
+
+def dice_loss(y_hat, y):
+    smooth = 1e-6
+    y_hat = y_hat.view(-1)
+    y = y.view(-1)
+    intersection = (y_hat * y).sum()
+    union = y_hat.sum() + y.sum()
+    dice = (2 * intersection + smooth) / (union + smooth)
+    return 1 - dice

src/vgg16_model.py

@@ -0,0 +1,209 @@
+import torch
+import torch.nn as nn
+import segmentation_models_pytorch as smp
+from torchmetrics import F1Score, Precision, Recall, JaccardIndex
+import pytorch_lightning as pl
+import wandb
+from torch.optim import Adam
+from torch.optim.lr_scheduler import StepLR
+
+class smp_model_vgg16(nn.Module):
+    def __init__(self, in_channels, out_channels, model_type, num_classes, encoder_weights):
+        super(smp_model_vgg16, self).__init__()
+        
+        # Ensure that model is always initialized
+        self.model = smp.Unet(
+            encoder_name='vgg16',
+            encoder_weights=encoder_weights,
+            in_channels=in_channels,  # The number of input channels, which is 14
+            classes=num_classes,       # Output classes, which is 1
+        )
+
+    def forward(self, x):
+        return self.model(x)
+
+class LandslideModel(pl.LightningModule):
+    def __init__(self, config, alpha=0.5):
+        super(LandslideModel, self).__init__()
+
+        model_type = config['model_config']['model_type']
+        in_channels = config['model_config']['in_channels']
+        num_classes = config['model_config']['num_classes']
+        self.alpha = alpha  # Assign the alpha value to the class variable
+        self.lr = config['train_config']['lr']
+
+        if model_type == 'unet':
+            self.model = UNet(in_channels=in_channels, out_channels=num_classes)
+        else:
+            encoder_weights = config['model_config']['encoder_weights']
+            # Use the custom smp_model_vgg16 instead of smp_model
+            self.model = smp_model_vgg16(in_channels=in_channels, 
+                                          out_channels=num_classes, 
+                                          model_type=model_type, 
+                                          num_classes=num_classes, 
+                                          encoder_weights=encoder_weights)
+
+        self.weights = torch.tensor([5], dtype=torch.float32).to(self.device)
+        self.wce = nn.BCELoss(weight=self.weights)
+
+        self.train_f1 = F1Score(task='binary')
+        self.val_f1 = F1Score(task='binary')
+
+        self.train_precision = Precision(task='binary')
+        self.val_precision = Precision(task='binary')
+
+        self.train_recall = Recall(task='binary')
+        self.val_recall = Recall(task='binary')
+
+        self.train_iou = JaccardIndex(task='binary')
+        self.val_iou = JaccardIndex(task='binary')
+
+    def forward(self, x):
+        return self.model(x)
+
+    def training_step(self, batch, batch_idx):
+        x, y = batch
+        y_hat = torch.sigmoid(self(x))
+
+        wce_loss = self.wce(y_hat, y)
+        dice = dice_loss(y_hat, y)
+
+        combined_loss = (1 - self.alpha) * wce_loss + self.alpha * dice
+
+        precision = self.train_precision(y_hat, y)
+        recall = self.train_recall(y_hat, y)
+        iou = self.train_iou(y_hat, y)
+        loss_f1 = self.train_f1(y_hat, y)
+
+        self.log('train_precision', precision)
+        self.log('train_recall', recall)
+        self.log('train_wce', wce_loss)
+        self.log('train_dice', dice)
+        self.log('train_iou', iou)
+        self.log('train_f1', loss_f1)
+        self.log('train_loss', combined_loss)
+        return {'loss': combined_loss}
+
+    def validation_step(self, batch, batch_idx):
+        x, y = batch
+        y_hat = torch.sigmoid(self(x))
+
+        wce_loss = self.wce(y_hat, y)
+        dice = dice_loss(y_hat, y)
+
+        combined_loss = (1 - self.alpha) * wce_loss + self.alpha * dice
+
+        precision = self.val_precision(y_hat, y)
+        recall = self.val_recall(y_hat, y)
+        iou = self.val_iou(y_hat, y)
+        loss_f1 = self.val_f1(y_hat, y)
+
+        self.log('val_precision', precision)
+        self.log('val_recall', recall)
+        self.log('val_wce', wce_loss)
+        self.log('val_dice', dice)
+        self.log('val_iou', iou)
+        self.log('val_f1', loss_f1)
+        self.log('val_loss', combined_loss)
+
+        if self.current_epoch % 10 == 0:
+            x = (x - x.min()) / (x.max() - x.min())
+            x = x[:, 0:3]
+            x = x.permute(0, 2, 3, 1)
+            y_hat = (y_hat > 0.5).float()
+
+            class_labels = {0: "no landslide", 1: "landslide"}  # Define class_labels here
+
+            self.logger.experiment.log({
+                "image": wandb.Image(x[0].cpu().detach().numpy(), masks={ 
+                    "predictions": {
+                        "mask_data": y_hat[0][0].cpu().detach().numpy(),
+                        "class_labels": class_labels
+                    },
+                    "ground_truth": {
+                        "mask_data": y[0][0].cpu().detach().numpy(),
+                        "class_labels": class_labels
+                    }
+                })
+            })
+        return {'val_loss': combined_loss}
+
+    def configure_optimizers(self):
+        optimizer = Adam(self.parameters(), lr=self.lr)
+        scheduler = StepLR(optimizer, step_size=30, gamma=0.1)
+        return [optimizer], [scheduler]
+
+class Block(nn.Module):
+    def __init__(self, inputs=3, middles=64, outs=64):
+        super().__init__()
+
+        self.conv1 = nn.Conv2d(inputs, middles, 3, 1, 1)
+        self.conv2 = nn.Conv2d(middles, outs, 3, 1, 1)
+        self.relu = nn.ReLU()
+        self.bn = nn.BatchNorm2d(outs)
+        self.pool = nn.MaxPool2d(2, 2)
+
+    def forward(self, x):
+        x = self.relu(self.conv1(x))
+        x = self.relu(self.bn(self.conv2(x)))
+        return self.pool(x), x
+
+class UNet(nn.Module):
+    def __init__(self, in_channels=3, out_channels=1):
+        super().__init__()
+
+        self.en1 = Block(in_channels, 64, 64)
+        self.en2 = Block(64, 128, 128)
+        self.en3 = Block(128, 256, 256)
+        self.en4 = Block(256, 512, 512)
+        self.en5 = Block(512, 1024, 512)
+
+        self.upsample4 = nn.ConvTranspose2d(512, 512, 2, stride=2)
+        self.de4 = Block(1024, 512, 256)
+
+        self.upsample3 = nn.ConvTranspose2d(256, 256, 2, stride=2)
+        self.de3 = Block(512, 256, 128)
+
+        self.upsample2 = nn.ConvTranspose2d(128, 128, 2, stride=2)
+        self.de2 = Block(256, 128, 64)
+
+        self.upsample1 = nn.ConvTranspose2d(64, 64, 2, stride=2)
+        self.de1 = Block(128, 64, 64)
+
+        self.conv_last = nn.Conv2d(64, out_channels, kernel_size=1, stride=1, padding=0)
+
+    def forward(self, x):
+        x, e1 = self.en1(x)
+        x, e2 = self.en2(x)
+        x, e3 = self.en3(x)
+        x, e4 = self.en4(x)
+        _, x = self.en5(x)
+
+        x = self.upsample4(x)
+        x = torch.cat([x, e4], dim=1)
+        _, x = self.de4(x)
+
+        x = self.upsample3(x)
+        x = torch.cat([x, e3], dim=1)
+        _, x = self.de3(x)
+
+        x = self.upsample2(x)
+        x = torch.cat([x, e2], dim=1)
+        _, x = self.de2(x)
+
+        x = self.upsample1(x)
+        x = torch.cat([x, e1], dim=1)
+        _, x = self.de1(x)
+
+        x = self.conv_last(x)
+
+        return x
+
+def dice_loss(y_hat, y):
+    smooth = 1e-6
+    y_hat = y_hat.view(-1)
+    y = y.view(-1)
+    intersection = (y_hat * y).sum()
+    union = y_hat.sum() + y.sum()
+    dice = (2 * intersection + smooth) / (union + smooth)
+    return 1 - dice