Upload folder using huggingface_hub

densenet/custom_loss.py

@@ -0,0 +1,61 @@
+import torch
+import torch.nn as nn
+
+
+class CombinedLoss(nn.Module):
+    """
+    Combined loss function that includes CrossEntropyLoss and Dice Loss.
+    The combined loss is a weighted sum of the two losses.
+    Args:
+        alpha (float): Weight for CrossEntropyLoss. The weight for Dice Loss is (1 - alpha).
+        smooth (float): Smoothing factor for Dice Loss to avoid division by zero.
+    """
+    def __init__(self, alpha=0.25, smooth=1e-8):  # alpha balances the two losses
+        super().__init__()
+        self.alpha = alpha
+        self.ce = nn.CrossEntropyLoss()
+        self.smooth = smooth
+
+    def forward(self, preds, targets):
+        loss_ce = self.ce(preds, targets)
+        loss_dice = 1-self.dice(preds, targets)
+        return self.alpha * loss_ce + (1 - self.alpha) * loss_dice
+
+    def dice_per_class(self, preds, targets):
+        """
+        This function computes the Dice score for each slide. And outputs
+        the average Dice score for all slides.
+        Args:
+            preds (torch.Tensor): The predicted mask of shape (B, H, W).
+            targets (torch.Tensor): The ground truth mask of shape (B, H, W).
+        Returns:
+            float: The average Dice score for all slides.
+        """
+        B, H, W = targets.shape
+        total_dice = 0
+        for i in range(B):
+            intersection = torch.sum(preds[i] * targets[i])
+            union = torch.sum(preds[i]) + torch.sum(targets[i])
+            dice = (2*intersection + self.smooth)/(union + self.smooth) 
+            total_dice += dice
+        return total_dice/B
+        
+
+    def dice(self, preds, targets):
+        """
+        This function computes the Dice score for each class. And outputs
+        the average Dice score for all classes.
+        Args:
+            preds (torch.Tensor): The predicted mask of shape (B, C, H, W).
+            targets (torch.Tensor): The ground truth mask of shape (B, C, H, W).
+        Returns:
+            float: The average Dice score for all classes.
+        """
+
+        B, C, H, W = targets.shape
+        dice_for_each_class = 0
+        
+        for i in range(C):
+            dice_for_each_class += self.dice_per_class(preds[:,i,:,:], targets[:,i,:,:])
+        return dice_for_each_class/C
+        

densenet/dense_block.py

@@ -0,0 +1,59 @@
+import torch
+import torch.nn as nn
+from layer import Layer
+
+
+
+class DenseBlock(nn.Module):
+    """
+    Dense block for DenseNet.
+    This block consists of multiple layers where each layer's output is concatenated
+    to the input of the next layer.
+    This class was developed following the paper:
+    "The One Hundred Layers Tiramisu: Fully Convolutional DenseNets for Semantic Segmentation"
+    and the reference paper of this project.
+    Args:
+        in_channels (int): Number of input channels.
+        num_layers (int): Number of layers in the dense block.
+        growth_rate (int): Growth rate for the dense block.
+    """
+    def __init__(self, in_channels, num_layers, growth_rate):
+        super(DenseBlock, self).__init__()
+        layers = []
+        for i in range(num_layers):
+            layers.append(Layer(in_channels + i * growth_rate, growth_rate))
+        self.block = nn.Sequential(*layers)
+    
+    def forward(self, x):
+        outputs = []
+        for layer in self.block:
+            output = layer(x)
+            outputs.append(output)
+            x = torch.cat([x, output], dim=1)  # Concatenate along channel axis
+        
+        # implementation from the model found in the paper:  https://arxiv.org/pdf/1611.09326
+        output = torch.cat(outputs, dim=1)  # Concatenate all outputs
+        return output
+
+
+class InceptionX(nn.Module):
+    """
+    InceptionX block with three branches of different kernel sizes.
+    This is the first block of the DenseNet model.
+    Args:
+        in_channels (int): Number of input channels.
+    """
+    def __init__(self, in_channels):
+        super(InceptionX, self).__init__()
+        # Each branch with a different padding to keep the size of the output
+        self.branch_3x3 = nn.Conv2d(in_channels, 16, kernel_size=3, padding=1, bias=False)
+        self.branch_5x5 = nn.Conv2d(in_channels, 4, kernel_size=5, padding=2, bias=False)
+        self.branch_7x7 = nn.Conv2d(in_channels, 4, kernel_size=7, padding=3, bias=False)
+        self.bn = nn.BatchNorm2d(24)
+
+    def forward(self, x):
+        out_3x3 = self.branch_3x3(x)
+        out_5x5 = self.branch_5x5(x)
+        out_7x7 = self.branch_7x7(x)
+        out = torch.cat([out_3x3, out_5x5, out_7x7], dim=1)  # concatenate along channel axis
+        return self.bn(out)

densenet/densenet.py

@@ -0,0 +1,128 @@
+# The denseNet model implementation in PyTorch is based on the paper:
+# Densely Connected Fully Convolutional Network for Short-Axis Cardiac Cine MR Image Segmentation and Heart Diagnosis Using Random Forest 
+# https://link.springer.com/chapter/10.1007/978-3-319-75541-0_15#Tab3
+import torch
+import torch.nn as nn
+import sys
+from scipy.ndimage import label
+import numpy as np
+
+sys.path.append("./densenet")  # Add the parent directory to the path
+from dense_block import DenseBlock, InceptionX
+from transitions import TransitionDown, TransitionUp
+
+
+class DenseNet(nn.Module):
+    def __init__(self):
+        """
+        This is the DenseNet model for image segmentation based on the paper:
+        "The One Hundred Layers Tiramisu: Fully Convolutional DenseNets for Semantic Segmentation"
+        and the reference paper of this project.
+        The layers are organized as follows:
+        - Inception_X
+        - Dense Block (3 layers)
+        - Transition Down
+        - Dense Block (4 layers)
+        - Transition Down
+        - Dense Block (5 layers)
+        - Transition Down
+        - Bottleneck
+        - Transition Up
+        - Dense Block (5 layers)
+        - Transition Up
+        - Dense Block (4 layers)
+        - Transition Up
+        - Dense Block (3 layers)
+        - 1x1 convolution
+        - softmax activation
+        """
+        super(DenseNet, self).__init__()
+        growth_rate = 8
+
+        self.inception=InceptionX(1) # output channels = 24
+        self.downdense1=DenseBlock(24, 3, growth_rate=growth_rate) # output channels = 24
+        self.td1=TransitionDown(48, 48)
+        self.downdense2=DenseBlock(48, 4, growth_rate=growth_rate) #output channels = 32
+        self.td2=TransitionDown(80, 80)
+        self.downdense3=DenseBlock(80, 5, growth_rate=growth_rate) # output channels = 40
+        self.td3=TransitionDown(120, 120)
+        self.bottleneck=DenseBlock(120, 8, growth_rate=7) # Bottleneck output channels = 56
+        self.tu1=TransitionUp(56, 56)
+        self.updense1=DenseBlock(176, 5, growth_rate=growth_rate) # output channels = 40
+        self.tu2=TransitionUp(40, 40)
+        self.updense2=DenseBlock(120, 4, growth_rate=growth_rate) # output channels = 32
+        self.tu3=TransitionUp(32, 32)
+        self.updense3=DenseBlock(80, 3, growth_rate=growth_rate) # output channels = 24
+        self.finalconv=nn.Conv2d(24, out_channels=4, kernel_size=1) # output channels = 4
+        # softmax activation
+        self.softmax = nn.Softmax(dim=1) # output channels = 4
+    
+    def forward(self, x):
+        x = self.inception(x) # size 128x128
+        x1 = self.downdense1(x)
+        x11 = torch.cat([x, x1], dim=1) # channels = 48
+        x12 = self.td1(x11) 
+        x2 = self.downdense2(x12)
+        x21 = torch.cat([x12, x2], dim=1) # channels = 56
+        x22 = self.td2(x21)
+        x3 = self.downdense3(x22)
+        x31 = torch.cat([x22, x3], dim=1) # channels = 120
+        x32 = self.td3(x31)
+        x4 = self.bottleneck(x32)
+        x42 = self.tu1(x4)
+        x43 = torch.cat([x31, x42], dim=1) 
+        x44 = self.updense1(x43)
+        x45 = self.tu2(x44)
+        x46 = torch.cat([x21, x45], dim=1)
+        x47 = self.updense2(x46)
+        x48 = self.tu3(x47)
+        x49 = torch.cat([x11, x48], dim=1)
+        x5 = self.updense3(x49)
+        x51 = self.finalconv(x5)
+        x52 = self.softmax(x51)
+        return x52
+        # NOTE: I´m aware that this code is a little messy with the name of the variables.
+        # However I did it by hand without LLM help, so I figure it would be nice to leave
+        # it like this to show that
+
+    def load_model(self, model_path):
+        """
+        Load the model weights from a file.
+        Args:
+            model_path (str): Path to the model weights file.
+        """
+        self.load_state_dict(torch.load(model_path, map_location=torch.device('cuda' if torch.cuda.is_available() else 'cpu')))
+    
+    def get_largest_component(self, mask):
+        """
+        This function takes a mask and returns the largest connected component of the mask.
+        
+        Args:
+            mask: A 3D mask (B,W,H)
+        REturns:
+            A 3D mask with only the largest connected component. (B,W,H)
+        """
+        if len(mask.shape) != 3:
+            raise ValueError("The input mask tensor must be a 3D mask.")
+        
+        output_mask = np.zeros_like(mask)
+        
+        for slide in range(mask.shape[0]):
+            img = mask[slide]
+            structure = [[1,1,1],[1,1,1],[1,1,1]]
+            labeled, num_features = label(img, structure=structure)
+            
+            if num_features == 0:
+                return mask  # No components found
+            
+            # Find the largest component
+            counts = np.bincount(labeled.flat)
+            counts[0] = 0  # Ignore background count
+            largest_label = counts.argmax()
+            
+            # Create mask for the largest component
+            largest_component = (labeled == largest_label)
+            output_mask[slide] = largest_component
+        
+        return output_mask
+        

densenet/layer.py

@@ -0,0 +1,23 @@
+import torch.nn as nn
+
+class Layer(nn.Module):
+    def __init__(self, in_channels, out_channels):
+        """
+        DenseNet layer with Batch Normalization, ELU activation, 
+        Convolution, and Dropout.
+        Args:
+            in_channels (int): Number of input channels.
+            out_channels (int): Number of output channels. This is the growth rate.
+        """
+        super(Layer, self).__init__()
+        self.block = nn.Sequential(
+            nn.BatchNorm2d(in_channels),
+            nn.ELU(inplace=True), # Exponential ReLU
+            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
+            nn.Dropout2d(p=0.2)
+        )
+    
+    def forward(self, x):
+        x = self.block(x)
+        return x
+

densenet/transitions.py

@@ -0,0 +1,54 @@
+import torch.nn as nn
+
+class TransitionDown(nn.Module):
+    """
+    Transition down block for DenseNet.
+    This is the downsampling used in the first half of the network.
+    The block downsamples the input by a factor of 2 using MaxPooling.
+    Args:
+        in_channels (int): Number of input channels.
+        out_channels (int): Number of output channels.
+    """
+    def __init__(self, in_channels, out_channels):
+        super(TransitionDown, self).__init__()
+        self.block = nn.Sequential(
+            nn.BatchNorm2d(in_channels),
+            nn.ELU(inplace=True), # Exponential ReLU
+            nn.Conv2d(in_channels, out_channels, kernel_size=1),
+            nn.Dropout2d(p=0.2),
+            nn.MaxPool2d(kernel_size=2, stride=2)  # Downsamples by 2
+        )
+    
+    def forward(self, x):
+        return self.block(x)
+
+
+class TransitionUp(nn.Module):
+    """
+    Transition up block for DenseNet.
+    This is the upsampling used in the second half of the network.
+    The block upsamples the input by a factor of 2 using ConvTranspose.
+    Args:
+        in_channels (int): Number of input channels.
+        out_channels (int): Number of output channels.
+    """
+    def __init__(self, in_channels, out_channels):
+        super(TransitionUp, self).__init__()
+        self.convtrans = nn.ConvTranspose2d(
+            in_channels, 
+            out_channels, 
+            kernel_size=3, 
+            stride=2,
+            padding=1,
+            # not extremely happy with this output padding
+            # but it has to be there because otherwise the 
+            # output size will be necesarily a odd number
+            # according to the formula
+            # Hout​=(Hin​−1)×stride[0]−2×padding[0]+dilation[0]×(kernel_size[0]−1)+output_padding[0]+1
+            # source: https://pytorch.org/docs/stable/generated/torch.nn.ConvTranspose2d.html#torch.nn.ConvTranspose2d
+            output_padding=1, 
+            )  # Upsamples by 2
+    
+    def forward(self, x):
+        return self.convtrans(x)
+