Upload 7 files

app.py

@@ -0,0 +1,80 @@
+import streamlit as st
+import torch
+from PIL import Image
+import torchvision.transforms as transforms
+from train5 import deeplabv3_encoder_decoder
+import numpy as np
+
+# Function to load the model
+def load_model(model_path):
+    model = deeplabv3_encoder_decoder()
+
+    try:
+        model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu')))
+        model.eval()
+        return model
+    except Exception as e:
+        st.error(f"Error loading model: {e}")
+        return None
+
+# Path to the model
+model_path = '/teamspace/studios/this_studio/Segmentation/model.pth'
+  
+
+# Load the trained model
+model = load_model(model_path)
+
+if model:
+    # Create a Streamlit app
+    st.title('Aerial Image Segmentation')
+
+    # Add a file uploader to the app
+    uploaded_file = st.file_uploader("Choose an image...", type="jpg")
+
+    if uploaded_file is not None:
+        image = Image.open(uploaded_file)
+
+        # Display the original image
+        st.image(image, caption='Uploaded Image.', use_column_width=True)
+
+        # Preprocess the image
+        data_transform = transforms.Compose([
+            transforms.Resize((512, 512)),
+            transforms.ToTensor()]
+        )
+        image = data_transform(image)
+        image = image.unsqueeze(0)  # add a batch dimension
+
+        # Pass the image through the model
+        with torch.no_grad():
+            output = model(image)
+
+        # Define the color map and class labels
+        color_map = {
+            0: np.array([255, 34, 133]),  # Unlabeled
+            1: np.array([0, 252, 199]),   # Early Blight
+            2: np.array([86, 0, 254]),    # Late Blight
+            3: np.array([0, 0, 0])        # Leaf Minor
+        }
+
+        class_labels = {
+            0: 'Unlabeled',
+            1: 'Early Blight',
+            2: 'Late Blight',
+            3: 'Leaf Minor'
+        }
+
+
+        for k, v in class_labels.items():
+            st.sidebar.markdown(f'<div style="color:rgb{tuple(color_map[k])};">{v}</div>', unsafe_allow_html=True)
+
+
+        output = torch.argmax(output.squeeze(), dim=0).detach().cpu().numpy()
+
+
+        output_rgb = np.zeros((output.shape[0], output.shape[1], 3), dtype=np.uint8)
+        for k, v in color_map.items():
+            output_rgb[output == k] = v
+
+
+        st.image(output_rgb, caption='Segmented Image.', use_column_width=True)

dataloader.py

@@ -0,0 +1,66 @@
+import os
+from torch.utils.data import Dataset
+from PIL import Image
+import torchvision.transforms as transforms
+import numpy as np
+
+class AerialImageDataset(Dataset):
+    def __init__(self, image_dir, mask_dir, transform=None):
+        self.image_dir = image_dir
+        self.mask_dir = mask_dir
+        self.transform = transform
+        self.images = os.listdir(self.image_dir)
+        self.Hex_Classes = [
+            ('Unlabeled', '#FF2285'),
+            ('Early Blight','#00FCC7'),
+            ('Late Blight', '#5600FE'),
+            ('Leaf Minor', '#000000')
+        ]
+
+    def __len__(self):
+        return len(self.images)
+   
+    def __getitem__(self, idx):
+        
+        img_path = os.path.join(self.image_dir, self.images[idx])
+        mask_path = os.path.join(self.mask_dir, self.images[idx].replace('.jpg', '.png'))
+
+        image = Image.open(img_path)
+        mask = Image.open(mask_path).resize((512, 512))
+        # print(mask.size)
+
+        # print(mask.size)
+        
+        mask = np.array(mask)
+        mask = self.encode_segmap(mask)
+        mask = mask.astype(np.uint8)  # Convert data type to uint8
+        # print(mask.shape)
+        mask = Image.fromarray(mask) # Convert mask -> PIL
+
+        if self.transform:
+            image = self.transform(image)
+            mask = self.transform(mask)
+
+        return image, mask
+
+    def encode_segmap(self, mask):
+        mask = mask.astype(int)
+        label_mask = np.zeros((mask.shape[0], mask.shape[1]), dtype=np.int16)  # height, width -> 0
+        for i, (name, color) in enumerate(self.Hex_Classes):
+            if mask.ndim == 3:
+                label_mask[(mask[:,:,0] == int(color[1:3], 16)) & (mask[:,:,1] == int(color[3:5], 16)) & (mask[:,:,2] == int(color[5:7], 16))] = i
+            elif mask.ndim == 2:
+                label_mask[(mask == int(color[1:3], 16))] = i
+                # print("Warning ndim = 2")
+                # return None
+        
+        msk = np.zeros((512,512,4))
+        for i in [0,1,2,3]:
+            if i == 0:
+                msk_ind = np.where(label_mask == i, 4, 0)
+                msk[:,:,i] = msk_ind
+            else:
+                msk_ind = np.where(label_mask == i, i, 0)
+                msk[:,:,i] = msk_ind
+        # print("mask shape",type(msk))
+        return msk

losses.py

@@ -0,0 +1,93 @@
+# from .train3 import deeplabv3_encoder_decoder
+# # from .train3 import pl
+# # from .train3 import torch
+# import torch.nn as nn
+# import torch.nn.functional as F
+# import torch
+# class mIoULoss(nn.Module):
+#     def __init__(self, weight=None, size_average=True, n_classes=4):
+#         super().__init__()
+#         self.classes = n_classes
+
+#     def to_one_hot(self, tensor):
+#         n, h, w = tensor.size()
+#         one_hot = torch.zeros(n, self.classes, h, w).to(tensor.device)
+#         one_hot.scatter_(1, tensor.unsqueeze(1), 1)
+#         return one_hot
+
+#     def forward(self, inputs, target):
+#         N = inputs.size(0)
+#         inputs = F.softmax(inputs, dim=1)
+#         target_oneHot = self.to_one_hot(target)
+#         inter = inputs * target_oneHot
+#         inter = inter.view(N, self.classes, -1).sum(2)
+#         union = inputs + target_oneHot - inter
+#         union = union.view(N, self.classes, -1).sum(2)
+#         loss = inter / union
+#         return 1 - loss.mean()
+
+import torch.nn as nn
+import torch.nn.functional as F
+import torch
+
+
+class DiceLoss(nn.Module):
+    def __init__(self, smooth=1.0):
+        super(DiceLoss, self).__init__()
+        self.smooth = smooth
+
+    def forward(self, preds, labels):
+        #
+        if preds.dim() == 4:
+            preds = torch.sigmoid(preds)
+
+        # Flatten the tensors
+        preds = preds.contiguous().view(-1)
+        labels = labels.contiguous().view(-1)
+
+        # Compute intersection and union
+        intersection = (preds * labels).sum()
+        dice = (2. * intersection + self.smooth) / (preds.sum() + labels.sum() + self.smooth)
+
+        # Dice loss is 1 - Dice coefficient
+        loss = 1 - dice
+        return loss
+    
+    
+class mIoULoss(nn.Module):
+    def __init__(self, weight=None, size_average=True, n_classes=4):  # Set n_classes to 4
+        super().__init__()
+        self.classes = n_classes
+
+    def to_one_hot(self, tensor):
+        tensor = tensor.long()  # Ensure tensor is a LongTensor
+        n, c, h, w = tensor.size()  # Adjust size extraction
+        one_hot = torch.zeros(n, self.classes, h, w).to(tensor.device)
+        one_hot.scatter_(1, tensor, 1)
+        return one_hot
+
+    def forward(self, inputs, target):
+        # inputs => N x Classes x H x W
+        # target_oneHot => N x Classes x H x W
+
+        N = inputs.size()[0]
+
+        # predicted probabilities for each pixel along channel
+        inputs = F.softmax(inputs, dim=1)
+        
+        # Numerator Product
+        target_oneHot = self.to_one_hot(target)
+        inter = inputs * target_oneHot
+        ## Sum over all pixels N x C x H x W => N x C
+        inter = inter.view(N, self.classes, -1).sum(2)
+
+        # Denominator 
+        union = inputs + target_oneHot - (inputs * target_oneHot)
+        ## Sum over all pixels N x C x H x W => N x C
+        union = union.view(N, self.classes, -1).sum(2)
+
+        loss = inter / union
+
+        ## Return average loss over classes and batch
+        return 1 - loss.mean()
+

main.py

@@ -0,0 +1,189 @@
+# # # import os
+# # # import pytorch_lightning as L
+# # # from dataloader import AerialImageDataset
+# # # from train5 import deeplabv3_encoder_decoder
+# # # from torch.utils.data import DataLoader
+# # # from torchvision.transforms import transforms
+# # # import torch
+
+# # # train_path = r"C:\Users\User\Downloads\Nishant\train"
+# # # val_path = r"C:\Users\User\Downloads\Nishant\val"
+
+# # # data_transform = transforms.Compose([
+# # #     transforms.Resize((512, 512)),
+# # #     transforms.ToTensor()
+# # # ])
+
+# # # train_dataset = AerialImageDataset(os.path.join(train_path, 'images'), os.path.join(train_path, 'masks'), transform=data_transform)
+# # # val_dataset = AerialImageDataset(os.path.join(val_path, 'images'), os.path.join(val_path, 'masks'), transform=data_transform)
+
+# # # train_loader = DataLoader(train_dataset, batch_size=2, shuffle=True)
+# # # val_loader = DataLoader(val_dataset, batch_size=2, shuffle=False)
+
+# # # model = deeplabv3_encoder_decoder()
+
+# # # # Adjust the refresh rate of the progress bar
+# # # trainer = L.Trainer(max_epochs=100, progress_bar_refresh_rate=20)  # Adjust the refresh rate as needed
+# # # trainer.fit(model, train_loader, val_loader)
+
+# # # torch.save(model.state_dict(), r"C:\Users\User\Downloads\Nishant\main.py\model.pth")
+
+# # import os
+# # import pytorch_lightning as pl
+# # from dataloader import AerialImageDataset
+# # from train5 import deeplabv3_encoder_decoder
+# # from torch.utils.data import DataLoader
+# # from torchvision.transforms import transforms
+# # import torch
+
+# # train_path = r"C:\Users\User\Downloads\Nishant\train"
+# # val_path = r"C:\Users\User\Downloads\Nishant\val"
+
+# # data_transform = transforms.Compose([
+# #     transforms.Resize((512, 512)),
+# #     transforms.ToTensor()
+# # ])
+
+# # train_dataset = AerialImageDataset(os.path.join(train_path, 'images'), os.path.join(train_path, 'masks'), transform=data_transform)
+# # val_dataset = AerialImageDataset(os.path.join(val_path, 'images'), os.path.join(val_path, 'masks'), transform=data_transform)
+
+# # train_loader = DataLoader(train_dataset, batch_size=2, shuffle=True)
+# # val_loader = DataLoader(val_dataset, batch_size=2, shuffle=False)
+
+# # model = deeplabv3_encoder_decoder()
+
+# # # Adjust other trainer parameters as needed
+# # trainer = pl.Trainer(max_epochs=100)
+# # trainer.fit(model, train_loader, val_loader)
+
+# # torch.save(model.state_dict(), r"C:\Users\User\Downloads\Nishant\main.py\model.pth")
+
+
+
+# #running code 
+# # import os
+# # import pytorch_lightning as pl
+# # from dataloader import AerialImageDataset
+# # from train5 import deeplabv3_encoder_decoder
+# # from torch.utils.data import DataLoader
+# # from torchvision.transforms import transforms
+# # import torch
+
+# # train_path = r"C:\Users\User\Downloads\Nishant\train"
+# # val_path = r"C:\Users\User\Downloads\Nishant\val"
+
+# # data_transform = transforms.Compose([
+# #     transforms.Resize((512, 512)),
+# #     transforms.ToTensor()
+# # ])
+
+# # train_dataset = AerialImageDataset(os.path.join(train_path, 'images'), os.path.join(train_path, 'masks'), transform=data_transform)
+# # val_dataset = AerialImageDataset(os.path.join(val_path, 'images'), os.path.join(val_path, 'masks'), transform=data_transform)
+
+# # train_loader = DataLoader(train_dataset, batch_size=2, shuffle=True)
+# # val_loader = DataLoader(val_dataset, batch_size=2, shuffle=False)
+
+# # model = deeplabv3_encoder_decoder()
+
+# # # Adjust other trainer parameters as needed
+# # trainer = pl.Trainer(num_sanity_val_steps=0, max_epochs=100)
+# # trainer.fit(model, train_loader, val_loader)
+
+# # torch.save(model.state_dict(), r"C:\Users\User\Downloads\Nishant\main.py\model.pth")
+
+# import os
+# import pytorch_lightning as pl
+# from dataloader import AerialImageDataset
+# from train5 import deeplabv3_encoder_decoder
+# from torch.utils.data import DataLoader
+# from torchvision.transforms import transforms
+# import torch
+# from pytorch_lightning.callbacks import ModelCheckpoint, EarlyStopping
+
+# train_path = r"C:\Users\User\Downloads\Nishant\train"
+# val_path = r"C:\Users\User\Downloads\Nishant\val"
+
+# data_transform = transforms.Compose([
+#     transforms.Resize((512, 512)),
+#     transforms.ToTensor()
+# ])
+
+# train_dataset = AerialImageDataset(os.path.join(train_path, 'images'), os.path.join(train_path, 'masks'), transform=data_transform)
+# val_dataset = AerialImageDataset(os.path.join(val_path, 'images'), os.path.join(val_path, 'masks'), transform=data_transform)
+
+# train_loader = DataLoader(train_dataset, batch_size=2, shuffle=True)
+# val_loader = DataLoader(val_dataset, batch_size=2, shuffle=False)
+
+# model = deeplabv3_encoder_decoder()
+
+
+# checkpoint_callback = ModelCheckpoint(
+#     monitor='val_loss',  
+#     dirpath='checkpoints',  
+#     filename='best_model',  
+#     save_top_k=1,  
+#     mode='min' 
+# )
+
+# early_stop_callback = EarlyStopping(
+#     monitor='val_loss',  
+#     patience=20,  
+#     verbose=True,
+#     mode='min'  
+# )
+
+
+# trainer = pl.Trainer(
+#     num_sanity_val_steps=0,
+#     max_epochs=100,
+#     callbacks=[checkpoint_callback, early_stop_callback]  # Pass both callbacks
+# )
+# trainer.fit(model, train_loader, val_loader)
+# torch.save(model.state_dict(), r"C:\Users\User\Downloads\Nishant\main.py\model.pth")
+import os
+import pytorch_lightning as pl
+from dataloader import AerialImageDataset
+from train5 import deeplabv3_encoder_decoder
+from torch.utils.data import DataLoader
+from torchvision.transforms import transforms
+import torch
+from pytorch_lightning.callbacks import ModelCheckpoint, EarlyStopping
+
+train_path = r"/teamspace/studios/this_studio/Segmentation/train"
+val_path = r"/teamspace/studios/this_studio/Segmentation/val"
+
+data_transform = transforms.Compose([
+    transforms.Resize((512, 512)),
+    transforms.ToTensor()
+])
+
+train_dataset = AerialImageDataset(os.path.join(train_path, 'images'), os.path.join(train_path, 'masks'), transform=data_transform)
+val_dataset = AerialImageDataset(os.path.join(val_path, 'images'), os.path.join(val_path, 'masks'), transform=data_transform)
+
+train_loader = DataLoader(train_dataset, batch_size=16, shuffle=True)
+val_loader = DataLoader(val_dataset, batch_size=16, shuffle=False)
+
+model = deeplabv3_encoder_decoder()
+
+checkpoint_callback = ModelCheckpoint(
+    monitor='val_loss',  
+    dirpath='checkpoints1',  
+    filename='best_model',  
+    save_top_k=1,  
+    mode='min'  # Save the model based on minimizing validation loss
+)
+
+early_stop_callback = EarlyStopping(
+    monitor='val_loss',  
+    patience=20,  
+    verbose=True,
+    mode='min'  
+)
+
+trainer = pl.Trainer(
+    num_sanity_val_steps=0,
+    max_epochs=1000,
+    callbacks=[checkpoint_callback, early_stop_callback]  # Pass both callbacks
+)
+trainer.fit(model, train_loader, val_loader)
+torch.save(model.state_dict(), r"/teamspace/studios/this_studio/Segmentation/model.pth")

model.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cd6ff7f738ef678fed4b0eb358462422d743004d7321b8378f12eae1f7fa93a9
+size 155201050

requirements.txt

@@ -0,0 +1,5 @@
+torch
+Pillow
+torchvision
+numpy
+pytorch-lightning==2.2.5

train5.py

@@ -0,0 +1,387 @@
+# import torch
+# import torch.nn as nn
+# import torch.nn.functional as F
+# import pytorch_lightning as pl
+# from losses import mIoULoss
+# from torchvision import models
+
+# class ASSP(nn.Module):
+#     def __init__(self, in_channels, out_channels=256, final_out_channels=4):
+#         super(ASSP, self).__init__()
+
+#         self.relu = nn.ReLU(inplace=True)
+
+#         # 1x1 convolution
+#         self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=1, padding=0, dilation=1, bias=False)
+#         self.bn1 = nn.BatchNorm2d(out_channels)
+
+#         # 3x3 convolutions with different dilation rates
+#         self.conv2 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=6, dilation=6, bias=False)
+#         self.bn2 = nn.BatchNorm2d(out_channels)
+
+#         self.conv3 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=12, dilation=12, bias=False)
+#         self.bn3 = nn.BatchNorm2d(out_channels)
+
+#         self.conv4 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=18, dilation=18, bias=False)
+#         self.bn4 = nn.BatchNorm2d(out_channels)
+
+#         # 1x1 convolution after global average pooling
+#         self.conv5 = nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0, dilation=1, bias=False)
+#         self.bn5 = nn.BatchNorm2d(out_channels)
+
+#         # Final 1x1 convolution to combine features
+#         self.convf = nn.Conv2d(out_channels * 5, final_out_channels, kernel_size=1, stride=1, padding=0, dilation=1, bias=False)
+#         self.bnf = nn.BatchNorm2d(final_out_channels)
+
+#         # Global average pooling
+#         self.adapool = nn.AdaptiveAvgPool2d(1)
+
+#     def forward(self, x):
+#         # 1x1 convolution
+#         x1 = self.conv1(x)
+#         x1 = self.bn1(x1)
+#         x1 = self.relu(x1)
+
+#         # 3x3 convolution with dilation 6
+#         x2 = self.conv2(x)
+#         x2 = self.bn2(x2)
+#         x2 = self.relu(x2)
+
+#         # 3x3 convolution with dilation 12
+#         x3 = self.conv3(x)
+#         x3 = self.bn3(x3)
+#         x3 = self.relu(x3)
+
+#         # 3x3 convolution with dilation 18
+#         x4 = self.conv4(x)
+#         x4 = self.bn4(x4)
+#         x4 = self.relu(x4)
+
+#         # Global average pooling, 1x1 convolution, and upsample
+#         x5 = self.adapool(x)
+#         x5 = self.conv5(x5)
+#         x5 = self.bn5(x5)
+#         x5 = self.relu(x5)
+#         x5 = F.interpolate(x5, size=x4.shape[-2:], mode='bilinear')
+
+#         # Concatenate all feature maps
+#         x = torch.cat((x1, x2, x3, x4, x5), dim=1)
+
+#         # Final 1x1 convolution
+#         x = self.convf(x)
+#         x = self.bnf(x)
+#         x = self.relu(x)
+
+#         return x
+    
+# class ResNet_50(nn.Module):
+#     def __init__(self, in_channels=3):  # Change default to 3 channels for RGB images
+#         super(ResNet_50, self).__init__()
+
+#         # Load the pre-trained ResNet-50 model
+#         self.resnet_50 = models.resnet50(weights='DEFAULT')
+
+#         # Modify the first convolutional layer to accept 3-channel input
+#         self.resnet_50.conv1 = nn.Conv2d(in_channels, 64, kernel_size=7, stride=2, padding=3, bias=False)
+
+#         # Use the layers up to the final layer before the fully connected layer
+#         self.resnet_50 = nn.Sequential(*list(self.resnet_50.children())[:-2])
+#         self.relu = nn.ReLU(inplace=True)
+
+#     def forward(self, x):
+#         x = self.resnet_50(x)
+#         return x
+
+# class deeplabv3_encoder_decoder(pl.LightningModule):
+#     def __init__(self, input_channels=3, output_channels=4):  # Use 4 channels for output
+#         super(deeplabv3_encoder_decoder, self).__init__()
+#         self.resnet = ResNet_50(in_channels=input_channels)
+#         self.aspp = ASSP(in_channels=2048, final_out_channels=4)
+#         self.conv = nn.Conv2d(in_channels=4, out_channels=output_channels, kernel_size=1)
+#         self.criterion = mIoULoss(n_classes=4)  # Set number of classes to 4
+
+#     def forward(self, x):
+#         _, _, h, w = x.shape
+#         x = self.resnet(x)  # Output should be [batch_size, 2048, H/32, W/32]
+#         x = self.aspp(x)
+#         x = F.interpolate(x, size=(h, w), mode='bilinear', align_corners=True)  # Upsample
+#         x = self.conv(x)  # Apply final convolution
+#         return x    
+
+#     def training_step(self, batch, batch_idx):
+#         images, masks = batch
+#         logits = self(images)
+#         loss = self.criterion(logits, masks)
+#         iou = calculate_iou(logits, masks)
+#         self.log('train_loss', loss)
+#         self.log('train_iou', iou)
+#         print(f'Training Loss: {loss}, IoU: {iou}')
+#         return loss
+
+#     def validation_step(self, batch, batch_idx):
+#         images, masks = batch
+#         logits = self(images)
+#         loss = self.criterion(logits, masks)
+#         iou = calculate_iou(logits, masks)
+#         self.log('val_loss', loss)
+#         self.log('val_iou', iou)
+#         print(f'Validation Loss: {loss}, IoU: {iou}')
+#         return loss
+
+#     def on_training_epoch_end(self, outputs):
+#         avg_iou = torch.stack([x['train_iou'] for x in outputs]).mean()
+#         self.log('avg_train_iou', avg_iou)
+#     def on_validation_epoch_end(self, outputs):
+#         avg_iou = torch.stack([x['val_iou'] for x in outputs]).mean()
+#         self.log('avg_val_iou', avg_iou)
+#     def configure_optimizers(self):
+#         optimizer = torch.optim.Adam(self.parameters(), lr=1e-3)
+#         return optimizer
+
+# def calculate_iou(logits, masks):
+#     # Calculate predictions from logits
+#     preds = torch.argmax(logits, dim=1)
+#     # Calculate intersection and union
+#     intersection = torch.sum(preds * masks)
+#     union = torch.sum((preds.bool() | masks.bool()).int())
+#     # Avoid division by zero
+#     iou = intersection / union if union != 0 else torch.tensor(0.0)
+#     return iou
+
+
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import pytorch_lightning as pl
+from losses import DiceLoss
+from torchvision import models
+import numpy as np
+import matplotlib.pyplot as plt
+
+class ASSP(nn.Module):
+    def __init__(self, in_channels, out_channels=256, final_out_channels=4):
+        super(ASSP, self).__init__()
+
+        self.relu = nn.ReLU(inplace=True)
+
+        # 1x1 convolution
+        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=1, padding=0, dilation=1, bias=False)
+        self.bn1 = nn.BatchNorm2d(out_channels)
+
+        # 3x3 convolutions with different dilation rates
+        self.conv2 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=6, dilation=6, bias=False)
+        self.bn2 = nn.BatchNorm2d(out_channels)
+
+        self.conv3 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=12, dilation=12, bias=False)
+        self.bn3 = nn.BatchNorm2d(out_channels)
+
+        self.conv4 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=18, dilation=18, bias=False)
+        self.bn4 = nn.BatchNorm2d(out_channels)
+
+        # 1x1 convolution after global average pooling
+        self.conv5 = nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0, dilation=1, bias=False)
+        self.bn5 = nn.BatchNorm2d(out_channels)
+
+        # Final 1x1 convolution to combine features
+        self.convf = nn.Conv2d(out_channels * 5, final_out_channels, kernel_size=1, stride=1, padding=0, dilation=1, bias=False)
+        self.bnf = nn.BatchNorm2d(final_out_channels)
+
+        # Global average pooling
+        self.adapool = nn.AdaptiveAvgPool2d(1)
+
+    def forward(self, x):
+        # 1x1 convolution
+        x1 = self.conv1(x)
+        x1 = self.bn1(x1)
+        x1 = self.relu(x1)
+
+        # 3x3 convolution with dilation 6
+        x2 = self.conv2(x)
+        x2 = self.bn2(x2)
+        x2 = self.relu(x2)
+
+        # 3x3 convolution with dilation 12
+        x3 = self.conv3(x)
+        x3 = self.bn3(x3)
+        x3 = self.relu(x3)
+
+        # 3x3 convolution with dilation 18
+        x4 = self.conv4(x)
+        x4 = self.bn4(x4)
+        x4 = self.relu(x4)
+
+        # Global average pooling, 1x1 convolution, and upsample
+        x5 = self.adapool(x)
+        x5 = self.conv5(x5)
+        x5 = self.bn5(x5)
+        x5 = self.relu(x5)
+        x5 = F.interpolate(x5, size=x4.shape[-2:], mode='bilinear')
+
+        # Concatenate all feature maps
+        x = torch.cat((x1, x2, x3, x4, x5), dim=1)
+
+        # Final 1x1 convolution
+        x = self.convf(x)
+        x = self.bnf(x)
+        x = self.relu(x)
+
+        return x
+    
+class ResNet_50(nn.Module):
+    def __init__(self, in_channels=3):  # Change default to 3 channels for RGB images
+        super(ResNet_50, self).__init__()
+
+        # Load the pre-trained ResNet-50 model
+        self.resnet_50 = models.resnet50(pretrained=True)
+
+        # Modify the first convolutional layer to accept 3-channel input
+        self.resnet_50.conv1 = nn.Conv2d(in_channels, 64, kernel_size=7, stride=2, padding=3, bias=False)
+
+        # Use the layers up to the final layer before the fully connected layer
+        self.resnet_50 = nn.Sequential(*list(self.resnet_50.children())[:-2])
+        self.relu = nn.ReLU(inplace=True)
+
+    def forward(self, x):
+        x = self.resnet_50(x)
+        return x
+
+# class deeplabv3_encoder_decoder(pl.LightningModule):
+#     def __init__(self, input_channels=3, output_channels=4):  # Use 4 channels for output
+#         super(deeplabv3_encoder_decoder, self).__init__()
+#         self.resnet = ResNet_50(in_channels=input_channels)
+#         self.aspp = ASSP(in_channels=2048, final_out_channels=4)
+#         self.conv = nn.Conv2d(in_channels=4, out_channels=output_channels, kernel_size=1)
+#         self.criterion = mIoULoss(n_classes=4)  # Set number of classes to 4
+
+#     def forward(self, x):
+#         _, _, h, w = x.shape
+#         x = self.resnet(x)  # Output should be [batch_size, 2048, H/32, W/32]
+#         x = self.aspp(x)
+#         x = F.interpolate(x, size=(h, w), mode='bilinear', align_corners=True)  # Upsample
+#         x = self.conv(x)  # Apply final convolution
+#         return x    
+
+#     def training_step(self, batch, batch_idx):
+#         images, masks = batch
+#         logits = self(images)
+#         loss = self.criterion(logits, masks)
+#         iou = calculate_iou(logits, masks)
+#         self.log('train_loss', loss)
+#         self.log('train_iou', iou)
+#         print(f'Training Loss: {loss}, IoU: {iou}')
+#         return loss
+
+#     def validation_step(self, batch, batch_idx):
+#         images, masks = batch
+#         logits = self(images)
+#         loss = self.criterion(logits, masks)
+#         iou = calculate_iou(logits, masks)
+#         self.log('val_loss', loss)
+#         self.log('val_iou', iou)
+#         print(f'Validation Loss: {loss}, IoU: {iou}')
+#         return loss
+
+#     def on_training_epoch_end(self, outputs):
+#         avg_iou = torch.stack([x['train_iou'] for x in outputs]).mean()
+#         self.log('avg_train_iou', avg_iou)
+
+#     def on_validation_epoch_end(self, outputs):
+#         avg_iou = torch.stack([x['val_iou'] for x in outputs]).mean()
+#         self.log('avg_val_iou', avg_iou)
+
+#     def configure_optimizers(self):
+#         optimizer = torch.optim.Adam(self.parameters(), lr=1e-3)
+#         return optimizer
+
+class deeplabv3_encoder_decoder(pl.LightningModule):
+    def __init__(self, input_channels=3, output_channels=4):  # Use 4 channels for output
+        super(deeplabv3_encoder_decoder, self).__init__()
+        self.resnet = ResNet_50(in_channels=input_channels)
+        self.aspp = ASSP(in_channels=2048, final_out_channels=4)
+        self.conv = nn.Conv2d(in_channels=4, out_channels=output_channels, kernel_size=1)
+        self.criterion = DiceLoss()  # Set number of classes to 4
+
+    def forward(self, x):
+        _, _, h, w = x.shape
+        x = self.resnet(x)  # Output should be [batch_size, 2048, H/32, W/32]
+        x = self.aspp(x)
+        x = F.interpolate(x, size=(h, w), mode='bilinear', align_corners=True)  # Upsample
+        x = self.conv(x)  # Apply final convolution
+        return x    
+
+    def training_step(self, batch, batch_idx):
+        images, masks = batch
+        logits = self(images)
+        loss = self.criterion(logits, masks)
+        # print("\n\n\n\n\n\n\n\n",masks.shape, logits.shape,"\n\n\n\n\n\n\n\n\n\n")
+        iou = compute_iou(logits, masks)
+        self.log('train_loss', loss)
+        self.log('train_iou', iou)
+        # print(f'Training Loss: {loss}, IoU: {iou}')
+        return loss
+
+    def validation_step(self, batch, batch_idx):
+        images, masks = batch
+        logits = self(images)
+        loss = self.criterion(logits, masks)
+        iou = compute_iou(logits, masks)
+        self.log('val_loss', loss)
+        self.log('val_iou', iou)
+        # print(f'Validation Loss: {loss}, IoU: {iou}')
+        return loss
+
+    def on_train_epoch_end(self):
+        avg_iou = self.trainer.callback_metrics['train_iou'].mean()
+        train_loss = self.trainer.logged_metrics.get('train_loss')
+        self.log('avg_train_iou', avg_iou)
+        print("avg train iou",avg_iou)
+        print("loss",train_loss)
+        # iou = calculate_iou(logits, masks)
+        # self.log('train_loss', loss)
+        # self.log('train_iou', iou)
+        # print(f'Training Loss: {loss}, IoU: {iou}')
+
+    def on_validation_epoch_end(self):
+        avg_iou = self.trainer.callback_metrics['val_iou'].mean()
+        val_loss = self.trainer.logged_metrics.get('val_loss')
+
+        self.log('avg_val_iou', avg_iou)
+        print("avg val iou",avg_iou)
+        print("val loss", val_loss) 
+
+    def configure_optimizers(self):
+        optimizer = torch.optim.Adam(self.parameters(), lr=1e-3)
+        return optimizer
+
+
+
+
+# def calculate_iou(logits, masks):
+#     # Calculate predictions from logits
+#     preds = torch.argmax(logits, dim=1)
+#     # Calculate intersection and union
+#     intersection = torch.sum(preds * masks)
+#     union = torch.sum((preds.bool() | masks.bool()).int())
+#     # Avoid division by zero
+#     iou = intersection / union if union != 0 else torch.tensor(0.0)
+#     return iou
+
+def compute_iou(preds,labels,threshold = 0.5 , epsilon = torch.finfo(torch.float).eps):
+    preds = torch.sigmoid(preds)
+    # print("preds shape",preds.shape)
+    preds = (preds>threshold).float()
+    # print("preds shape123",preds.shape)
+    # print("masks shape123",labels.shape)
+    # print("masks shape123",np.unique(labels.cpu().numpy()))
+    # plt.imshow(labels[0,:,:,:].T.cpu().numpy())
+    # plt.show()
+    n_classes = preds.shape[1]
+    iou_per_class = []
+    for i in range(n_classes):
+        intersection = (preds[:,i,:,:] * labels[:,i,:,:]).sum((1,2))
+        union = (preds[:,i,:,:]+ labels[:,i,:,:]).sum((1,2)) - intersection
+        iou = (intersection + epsilon) / (union + epsilon)
+        iou_per_class.append(iou.mean())
+    iou_mean = sum(iou_per_class)/ n_classes
+    return iou_mean