Upload model.py

model.py

@@ -0,0 +1,198 @@
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torch.utils.data import DataLoader, Dataset
+from torch.optim.lr_scheduler import ReduceLROnPlateau
+import numpy as np
+from PIL import Image
+import random
+import torch.nn.functional as F
+
+
+class CustomDataset(Dataset):
+
+    def __init__(self, red_dir, green_dir, blue_dir, nir_dir, mask_dir, pytorch=True):
+        super().__init__()
+        self.red_dir = red_dir
+        self.green_dir = green_dir
+        self.blue_dir = blue_dir
+        self.nir_dir = nir_dir
+        self.mask_dir = mask_dir
+
+        red_files = [f for f in self.red_dir.iterdir() if f.is_file()]
+        self.files = [self.combine_files(f) for f in red_files]
+        self.pytorch = pytorch
+
+
+    def combine_files(self, red_files: Path):
+        base_name = red_files.name
+
+        files = {
+            'red': red_files,
+            'green': self.green_dir / base_name.replace('red', 'green'),
+            'blue': self.blue_dir / base_name.replace('red', 'blue'),
+            'nir': self.nir_dir / base_name.replace('red', 'nir'),
+            'mask': self.mask_dir / base_name.replace('red', 'gt'),
+        }
+
+        for key, path in files.items():
+            if not path.exists():
+                raise FileNotFoundError(f'Missing file: {path} for {red_files}')
+        return files
+
+
+    def __len__(self):
+        return len(self.files)
+
+
+    def open_as_array(self, idx, invert=False, nir_included=False):
+        rgb = np.stack([
+                np.array(Image.open(self.files[idx]['red'])),
+                np.array(Image.open(self.files[idx]['green'])),
+                np.array(Image.open(self.files[idx]['blue']))
+            ], axis=2)
+
+        if nir_included:
+            nir = np.array(Image.open(self.files[idx]['nir']))
+            nir = np.expand_dims(nir, 2)
+            rgb = np.concatenate([rgb, nir], axis=2)
+
+        if invert:
+            rgb = rgb.transpose((2, 0, 1))
+
+        raw_rgb = (rgb / np.iinfo(rgb.dtype).max)
+        return raw_rgb
+
+    def open_mask(self,idx, expand_dims=True):
+        raw_mask = np.array(Image.open(self.files[idx]['mask']))
+        raw_mask = np.where(raw_mask == 255, 1, 0) # Transform the mask into binary array where pixels with value 256(white) become 1(clouds), pixels with 0 or anything else becomes 0(not clouds)
+
+        return np.expand_dims(raw_mask, 0) if expand_dims else raw_mask
+
+    def __getitem__(self, idx):
+        X = torch.tensor(self.open_as_array(idx, invert=True, nir_included=True), dtype=torch.float32)
+        y = torch.tensor(self.open_mask(idx, expand_dims=True), dtype=torch.float32)
+
+        return X, y
+
+
+class doubleConv(nn.Module):
+
+    def __init__(self, in_channels, out_channels):
+        super().__init__()
+
+        self.double_conv = nn.Sequential(
+            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
+            nn.BatchNorm2d(out_channels),
+            nn.ReLU(),
+            nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),
+            nn.BatchNorm2d(out_channels),
+            nn.ReLU()
+        )
+
+    def forward(self, x):
+        return self.double_conv(x)
+
+
+class downSample(nn.Module):
+
+    def __init__(self, in_channels, out_channels):
+        super().__init__()
+
+        self.conv = doubleConv(in_channels, out_channels)
+        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
+
+    def forward(self, x):
+        down = self.conv(x)
+        p = self.pool(down)
+
+        return down, p
+
+
+class upSample(nn.Module):
+    def __init__(self, in_channels, out_channels):
+        super().__init__()
+
+        self.up = nn.ConvTranspose2d(in_channels, out_channels, kernel_size=2, stride=2)
+        self.conv = doubleConv(out_channels * 2, out_channels)
+
+    def forward(self, x1, x2):
+        x1 = self.up(x1)
+        x = torch.cat([x1, x2], 1)
+        return self.conv(x)
+
+
+class SpatialAttention(nn.Module):
+
+    def __init__(self):
+        super().__init__()
+        self.conv = nn.Conv2d(in_channels=2, out_channels=1, kernel_size=3, padding=1)
+        self.sigmoid = nn.Sigmoid()
+
+    def forward(self, x):
+        avg_pooling = torch.mean(x, dim=1, keepdim=True)
+        max_pooling = torch.max(x, dim=1, keepdim=True)[0] # return on max values and not their indices
+        concat = torch.cat([avg_pooling, max_pooling], dim=1)
+        attention = self.conv(concat)
+        attention = self.sigmoid(attention)
+        output = x * attention
+        return output
+
+
+class UNet(nn.Module):
+    def __init__(self, in_channels, num_classes):
+        super().__init__()
+
+        self.down_conv1 = downSample(in_channels, 32)
+        self.down_conv2 = downSample(32, 64)
+        self.down_conv3 = downSample(64, 128)
+
+        self.bottleneck = doubleConv(128, 256)
+        self.spatial_attention = SpatialAttention()
+
+        self.up_conv1 = upSample(256, 128)
+        self.up_conv2 = upSample(128, 64)
+        self.up_conv3 = upSample(64, 32)
+
+        self.out = nn.Conv2d(in_channels=32 , out_channels=num_classes, kernel_size=1)
+
+    def forward(self, x):
+
+        down1, p1 = self.down_conv1(x)
+        down2, p2 = self.down_conv2(p1)
+        down3, p3 = self.down_conv3(p2)
+
+        b = self.bottleneck(p3)
+        b = self.spatial_attention(b)
+
+        up1 = self.up_conv1(b, down3)
+        up2 = self.up_conv2(up1, down2)
+        up3 = self.up_conv3(up2, down1)
+
+        output = self.out(up3)
+        return output
+
+def acc_fn(predb, yb):
+    preds = torch.sigmoid(predb)  # Convert logits to probabilities
+    preds = (preds > 0.5).float()  # Threshold at 0.5
+    return (preds == yb).float().mean()  # Compare with ground truth
+
+def calculate_metrics(y_true, y_pred):
+    TP = torch.sum((y_true == 1) & (y_pred == 1)).float()
+    TN = torch.sum((y_true == 0) & (y_pred == 0)).float()
+    FP = torch.sum((y_true == 0) & (y_pred == 1)).float()
+    FN = torch.sum((y_true == 1) & (y_pred == 0)).float()
+
+    jaccard = TP / (TP + FN + FP + 1e-10)
+    precision = TP / (TP + FP + 1e-10)
+    recall = TP / (TP + FN + 1e-10)
+    specificity = TN / (TN + FP + 1e-10)
+    overall_acc = (TP + TN) / (TP + TN + FP + FN + 1e-10)
+
+    return {
+        "Jaccard index": jaccard.item(),
+        "Precision": precision.item(),
+        "Recall": recall.item(),
+        "Specificity": specificity.item(),
+        "Overall Accuracy": overall_acc.item()
+    }