added docstrings for various scripts

app.py

@@ -19,6 +19,25 @@ from training.logger_utils import load_dict_from_json
 from training.dataset import get_dataloader_for_inference
 
 def run_inference(image_array, file_weights, num_classes=5, file_stats_json="training/image_stats.json"):
+    """
+    ---------
+    Arguments
+    ---------
+    image_array : ndarray
+        a numpy array of the image
+    file_weights : str
+        full path to weights file
+    num_classes : int
+        number of classes in the dataset
+    file_stats_json : str
+        full path to the json stats file for preprocessing
+
+    -------
+    Returns
+    -------
+    pred_mask_arr : ndarray
+        a numpy array of the prediction mask
+    """
     oil_spill_seg_model = ResNet50DeepLabV3Plus(
         num_classes=num_classes, pretrained=True
     )

training/dataset.py

@@ -13,6 +13,21 @@ from logger_utils import load_dict_from_json
 
 class M4DSAROilSpillDataset(Dataset):
     def __init__(self, dir_data, list_images, which_set="train", file_stats_json="image_stats.json"):
+        """
+        M4DSAROilSpillDataset class to load satellite image dataset
+
+        ----------
+        Attributes
+        ----------
+        dir_data : str
+            valid full directory path of the dataset
+        list_images : list
+            list of images in the directory
+        which_set : str
+            string indicates which set to be loaded (options = ["train", "test"])
+        file_stats_json : str
+            json file with image stats
+        """
         self.dir_data = dir_data
         self.which_set = which_set
         self.file_stats_json = file_stats_json
@@ -56,9 +71,29 @@ class M4DSAROilSpillDataset(Dataset):
 
 
     def __len__(self):
+        """
+        -------
+        Returns
+        -------
+        length : int
+            number of images in the dataset list
+        """
         return len(self._list_images)
 
     def __getitem__(self, idx):
+        """
+        ---------
+        Arguments
+        ---------
+        idx : int
+            index of the file
+
+        -------
+        Returns
+        -------
+        (image, label) : tuple of torch tensors
+            tuple of normalized image and label torch tensors
+        """
         file_image = os.path.join(self._dir_images, self._list_images[idx])
         file_label = os.path.join(self._dir_labels, self._list_labels[idx])
 
@@ -94,6 +129,25 @@ class M4DSAROilSpillDataset(Dataset):
         return image, label
 
 def get_dataloaders_for_training(dir_dataset, batch_size, random_state=None, num_workers=4):
+    """
+    ---------
+    Arguments
+    ---------
+    dir_dataset : str
+        full path to dataset directory
+    batch_size : int
+        batch size to be used
+    random_state : int
+        random state to be used for train / validation set split (default: None)
+    num_workers : int
+        number of workers to be used for dataloader (default: 4)
+
+    -------
+    Returns
+    -------
+    (train_dataset_loader, valid_dataset_loader) : tuple
+        tuple of torch dataloaders
+    """
     list_images = sorted(
         [f for f in os.listdir(os.path.join(dir_dataset, "train", "images")) if f.endswith(".jpg")]
     )
@@ -130,6 +184,23 @@ def get_dataloaders_for_training(dir_dataset, batch_size, random_state=None, num
     return train_dataset_loader, valid_dataset_loader
 
 def get_dataloader_for_inference(dir_dataset, batch_size=1, num_workers=4):
+    """
+    ---------
+    Arguments
+    ---------
+    dir_dataset : str
+        full path to dataset directory
+    batch_size : int
+        batch size to be used (default: 1)
+    num_workers : int
+        number of workers to be used for dataloader (default: 4)
+
+    -------
+    Returns
+    -------
+    (inference_dataset_loader, list_inference_images) : tuple
+        tuple of torch dataloader and a list of inference images
+    """
     list_inference_images = sorted(
         [f for f in os.listdir(os.path.join(dir_dataset, "test", "images")) if f.endswith(".jpg")]
     )

training/decoder_models.py

@@ -3,6 +3,23 @@ from torch import nn
 from torch.nn import functional as F
 
 class DeepLabV3(nn.Module):
+    """
+    DeepLabV3 class to build the DeepLabV3 decoder model
+
+    ----------
+    Attributes
+    ----------
+    in_channels : int
+        number of input channels to decoder model from the encoder model's output
+    num_classes : int
+        number of classes for which the decoder needs to be built
+    aspp_out_channels : int
+        number of output channels of the ASPP layer (default: 256)
+    final_out_channels : int
+        number of output channels before applying classification conv layer (default: 256)
+    aspp_dilate: list
+        a list of dilation rates to be used for conv layers in ASPP block (default: [12, 24, 36])
+    """
     def __init__(self, in_channels, num_classes, aspp_out_channels=256, final_out_channels=256, aspp_dilate=[12, 24, 36]):
         super().__init__()
         self.aspp_block = ASPPBlock(in_channels, aspp_dilate, aspp_out_channels=aspp_out_channels)
@@ -27,11 +44,45 @@ class DeepLabV3(nn.Module):
         return
 
     def forward(self, encoded_features):
+        """
+        ---------
+        Arguments
+        ---------
+        encoded_features : torch tensor
+            a tensor of encoded features from the encoder
+
+        -------
+        Returns
+        -------
+        final_output_feature : torch tensor
+            a tensor of final output logits
+        """
         aspp_output_feature = self.aspp_block(encoded_features)
         final_output_feature = self.classifier_conv_block(aspp_output_feature)
         return final_output_feature
 
 class DeepLabV3Plus(nn.Module):
+    """
+    DeepLabV3Plus class to build the DeepLabV3+ decoder model
+
+    ----------
+    Attributes
+    ----------
+    in_channels : int
+        number of input channels to decoder model from the encoder model's output
+    encoder_channels : int
+        number of channels from the intermediate layer of the encoder for merging
+    num_classes : int
+        number of classes for which the decoder needs to be built
+    encoder_projection_channels : int
+        number of resulting projection channels from the intermediate layer of the encoder for merging (default: 48)
+    aspp_out_channels : int
+        number of output channels of the ASPP layer (default: 256)
+    final_out_channels : int
+        number of output channels before applying classification conv layer (default: 256)
+    aspp_dilate: list
+        a list of dilation rates to be used for conv layers in ASPP block (default: [12, 24, 36])
+    """
     def __init__(self, in_channels, encoder_channels, num_classes, encoder_projection_channels=48,
         aspp_out_channels=256, final_out_channels=256, aspp_dilate=[12, 24, 36]):
 
@@ -66,6 +117,21 @@ class DeepLabV3Plus(nn.Module):
         return
 
     def forward(self, encoded_features, block_1_features):
+        """
+        ---------
+        Arguments
+        ---------
+        encoded_features : torch tensor
+            a tensor of encoded features from the encoder
+        block_1_features : torch tensor
+            a tensor of features from the intermediate layer from the encoder
+
+        -------
+        Returns
+        -------
+        final_output_feature : torch tensor
+            a tensor of final output logits
+        """
         encoder_connection = self.projection_conv(block_1_features)
         aspp_output_feature = self.aspp_block(encoded_features)
         aspp_output_feature = F.interpolate(
@@ -78,6 +144,19 @@ class DeepLabV3Plus(nn.Module):
         return final_output_feature
 
 class ASPPConvLayer(nn.Sequential):
+    """
+    ASPPConvLayer class to build the ASPPConvLayer used in ASPPBlock
+
+    ----------
+    Attributes
+    ----------
+    in_channels : int
+        number of input channels to ASPPConvLayer
+    out_channels : int
+        number of output channels from ASPPConvLayer
+    dilation : int
+        dilation rate
+    """
     def __init__(self, in_channels, out_channels, dilation):
         super().__init__()
         self.conv_block = nn.Sequential(
@@ -100,11 +179,35 @@ class ASPPConvLayer(nn.Sequential):
         return
 
     def forward(self, x):
+        """
+        ---------
+        Arguments
+        ---------
+        x : torch tensor
+            a tensor of input features
+
+        -------
+        Returns
+        -------
+        x : torch tensor
+            output of the ASPPConvLayer
+        """
         x = self.conv_block(x)
         return x
 
 class ASPPPoolingLayer(nn.Sequential):
     def __init__(self, in_channels, out_channels):
+        """
+        ASPPPoolingLayer class to build the ASPPPoolingLayer used in ASPPBlock
+
+        ----------
+        Attributes
+        ----------
+        in_channels : int
+            number of input channels to ASPPPoolingLayer
+        out_channels : int
+            number of output channels from ASPPPoolingLayer
+        """
         super().__init__()
         self.avg_pool_block = nn.Sequential(
             nn.AdaptiveAvgPool2d(1),
@@ -124,6 +227,19 @@ class ASPPPoolingLayer(nn.Sequential):
         return
 
     def forward(self, x):
+        """
+        ---------
+        Arguments
+        ---------
+        x : torch tensor
+            a tensor of input features
+
+        -------
+        Returns
+        -------
+        x : torch tensor
+            output of the ASPPPoolingLayer
+        """
         size = x.shape[2:]
         x = self.avg_pool_block(x)
         x = F.interpolate(x, size=size, mode="bilinear", align_corners=False)
@@ -131,6 +247,19 @@ class ASPPPoolingLayer(nn.Sequential):
 
 class ASPPBlock(nn.Module):
     def __init__(self, in_channels, atrous_rates, aspp_out_channels=256):
+        """
+        ASPPBlock class to build the ASPPBlock
+
+        ---------
+        Attributes
+        ----------
+        in_channels : int
+            number of input channels to ASPPBlock
+        atrous_rates : list
+            list of dilation rates
+        aspp_out_channels : int
+            number of output channels of the ASPPBlock
+        """
         super().__init__()
 
         self.aspp_init_conv = nn.Sequential(
@@ -168,6 +297,19 @@ class ASPPBlock(nn.Module):
         return
 
     def forward(self, x):
+        """
+        ---------
+        Arguments
+        ---------
+        x : torch tensor
+            a tensor of input features
+
+        -------
+        Returns
+        -------
+        x : torch tensor
+            output of the ASPPBlock
+        """
         aspp_outputs = []
         for aspp_layer in self.aspp_module_layers:
             aspp_outputs.append(aspp_layer(x))

training/encoder_models.py

@@ -27,6 +27,29 @@ class CustomResNet(nn.Module):
         replace_stride_with_dilation=None,
         norm_layer=None,
     ):
+        """
+        CustomResNet class to build the CustomResNet encoder model
+
+        ----------
+        Attributes
+        ----------
+        layers : list
+            list of number of layers in each residual block
+        block : object of block type
+            type of the residual block (options = [BasicBlock, Bottleneck])
+        zero_init_residual : bool
+            to indicate whether to use zero weights for BN
+        groups : int
+            indicates the number of groups (default: 1)
+        num_classes : int
+            indicates the number of classes (default: 1000)
+        width_per_group : int
+            indicates the width per group (default: 64)
+        replace_stride_with_dilation : list
+            a list indicating whether to replace stride with dilation (default: None)
+        norm_layer : object
+            object of type batch norm (default: None)
+        """
 
         super(CustomResNet, self).__init__()
 
@@ -120,6 +143,19 @@ class CustomResNet(nn.Module):
         return nn.Sequential(*layers)
 
     def forward(self, x):
+        """
+        ---------
+        Arguments
+        ---------
+        x : torch tensor
+            a tensor of input features
+
+        -------
+        Returns
+        -------
+        x : torch tensor
+            output of the CustomResNet
+        """
         x = self.conv1(x)
         x = self.bn1(x)
         x = self.relu(x)
@@ -135,6 +171,25 @@ class CustomResNet(nn.Module):
         return x
 
 def _resnet(block_type, layers, weights=None, progress=True):
+    """
+    ---------
+    Arguments
+    ---------
+    block_type : object
+        object of type block
+    layers : list
+        list of layers in each residual block
+    weights : object
+        object of type ResNet weights
+    progress : bool
+        indicates whether to show progress or not
+
+    -------
+    Returns
+    -------
+    model : object
+        model object of type CustomResNet
+    """
     model = CustomResNet(layers, block_type)
 
     if weights is not None:
@@ -146,8 +201,11 @@ def resnet18(pretrained=True):
     r"""ResNet-18 model from
     `"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_
 
-    Args:
-        pretrained (bool): If True, returns a model pre-trained on ImageNet
+    ---------
+    Arguments
+    ---------
+    pretrained : bool
+        if True, returns a model pre-trained on ImageNet
     """
     if pretrained:
         weights = ResNet18_Weights.IMAGENET1K_V1
@@ -159,8 +217,11 @@ def resnet34(pretrained=True):
     r"""ResNet-34 model from
     `"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_
 
-    Args:
-        pretrained (bool): If True, returns a model pre-trained on ImageNet
+    ---------
+    Arguments
+    ---------
+    pretrained : bool
+        if True, returns a model pre-trained on ImageNet
     """
     if pretrained:
         weights = ResNet34_Weights.IMAGENET1K_V1
@@ -172,8 +233,11 @@ def resnet50(pretrained=True):
     r"""ResNet-50 model from
     `"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_
 
-    Args:
-        pretrained (bool): If True, returns a model pre-trained on ImageNet
+    ---------
+    Arguments
+    ---------
+    pretrained : bool
+        if True, returns a model pre-trained on ImageNet
     """
     if pretrained:
         weights = ResNet50_Weights.IMAGENET1K_V1
@@ -186,8 +250,11 @@ def resnet101(pretrained=True):
     r"""ResNet-101 model from
     `"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_
 
-    Args:
-        pretrained (bool): If True, returns a model pre-trained on ImageNet
+    ---------
+    Arguments
+    ---------
+    pretrained : bool
+        if True, returns a model pre-trained on ImageNet
     """
     if pretrained:
         weights = ResNet101_Weights.IMAGENET1K_V1
@@ -209,13 +276,22 @@ class CustomEfficientNet(nn.Module):
     ):
         """
         EfficientNet V1 and V2 main class
-        Args:
-            inverted_residual_setting (Sequence[Union[MBConvConfig, FusedMBConvConfig]]): Network structure
-            dropout (float): The droupout probability
-            stochastic_depth_prob (float): The stochastic depth probability
-            num_classes (int): Number of classes
-            norm_layer (Optional[Callable[..., nn.Module]]): Module specifying the normalization layer to use
-            last_channel (int): The number of channels on the penultimate layer
+
+        ----------
+        Attributes
+        ----------
+            inverted_residual_setting : Sequence
+                network structure
+            dropout : float
+                the droupout probability
+            stochastic_depth_prob : float
+                the stochastic depth probability
+            num_classes : int
+                number of classes
+            norm_layer : object
+                object of type Module specifying the normalization layer to use
+            last_channel : int
+                the number of channels on the penultimate layer
         """
         super().__init__()
         self.dict_encoder_features = {}
@@ -321,6 +397,23 @@ def _efficientnet(
     progress=True,
     **kwargs: Any,
     ):
+    """
+    ---------
+    Arguments
+    ---------
+        inverted_residual_setting : Sequence
+            network structure
+        dropout : float
+            the droupout probability
+        last_channel : last_channel
+            the last channel
+        weights : object
+            object of type efficient_net weights
+        norm_layer : object
+            object of type Module specifying the normalization layer to use
+        progress : bool
+            indicates whether to show progress or not
+    """
     model = CustomEfficientNet(
         inverted_residual_setting,
         dropout,
@@ -335,6 +428,15 @@ def _efficientnet(
     return model
 
 def efficientnet_v2_s(pretrained=True, **kwargs: Any):
+    """
+    ---------
+    Arguments
+    ---------
+    pretrained : bool
+        if True, returns a model pre-trained on ImageNet
+    **kwargs :
+        additional arguments
+    """
     which_efficientnet = "efficientnet_v2_s"
     inverted_residual_setting, last_channel = _efficientnet_conf(which_efficientnet)
     if pretrained:
@@ -350,6 +452,15 @@ def efficientnet_v2_s(pretrained=True, **kwargs: Any):
     )
 
 def efficientnet_v2_m(pretrained=True, **kwargs: Any):
+    """
+    ---------
+    Arguments
+    ---------
+    pretrained : bool
+        if True, returns a model pre-trained on ImageNet
+    **kwargs :
+        additional arguments
+    """
     which_efficientnet = "efficientnet_v2_m"
     inverted_residual_setting, last_channel = _efficientnet_conf(which_efficientnet)
     if pretrained:
@@ -365,6 +476,15 @@ def efficientnet_v2_m(pretrained=True, **kwargs: Any):
     )
 
 def efficientnet_v2_l(pretrained=True, **kwargs: Any):
+    """
+    ---------
+    Arguments
+    ---------
+    pretrained : bool
+        if True, returns a model pre-trained on ImageNet
+    **kwargs :
+        additional arguments
+    """
     which_efficientnet = "efficientnet_v2_l"
     inverted_residual_setting, last_channel = _efficientnet_conf(which_efficientnet)
     if pretrained:

training/image_preprocessing.py

@@ -5,6 +5,25 @@ from skimage.io import imread
 
 class ImagePadder:
     def __init__(self, dir_images, pad_left=15, pad_right=15, pad_top=11, pad_bottom=11, file_anchor_image="img_0814.jpg"):
+        """
+        ImagePadder class for padding images
+
+        ----------
+        Attributes
+        ----------
+        dir_images : str
+            full directory path containing images
+        pad_left : int
+            number of pixels to be padded to the left of the input image (default: 15)
+        pad_right : int
+            number of pixels to be padded to the right of the input image (default: 15)
+        pad_top : int
+            number of pixels to be padded to the top of the input image (default: 11)
+        pad_bottom : int
+            number of pixels to be padded to the bottom of the input image (default: 11)
+        file_anchor_image : str
+            file with anchor image whose pixels will be used as a reference for padding (default: "img_0814.jpg")
+        """
         self._anchor_image = imread(os.path.join(dir_images, file_anchor_image))
         self._anchor_image_shape = self._anchor_image.shape
         self._pad_left = pad_left

training/inference.py

@@ -15,12 +15,41 @@ from seg_models import *
 from dataset import get_dataloader_for_inference
 
 def create_directory(dir_path):
+    """
+    ---------
+    Arguments
+    ---------
+    dir_path : str
+        full directory path that needs to be created if it does not exist
+    """
     if not os.path.isdir(dir_path):
         os.makedirs(dir_path)
         print(f"Created directory: {dir_path}")
     return
 
 def inference_loop(dataset_loader, list_images, model, dir_labels, dir_masks, num_classes, device, image_format=".png"):
+    """
+    ---------
+    Arguments
+    ---------
+    dataset_loader : object
+        object of type dataloader
+    list_images : list
+        list of images for which the inference needs to be run
+    model : object
+        object of type model
+    dir_labels : str
+        full directory path to save prediction labels
+    dir_masks : str
+        full directory path to save prediction masks
+    num_classes : int
+        number of classes in the dataset
+    device : str
+        device on which inference needs to be run
+    image_format : str
+        the extension format of the images (default: ".png")
+    """
+
     #  for lossless, always save labels and masks as png and not as jpeg
     model.eval()
     size = len(dataset_loader.dataset)

training/logger_utils.py

@@ -36,7 +36,16 @@ def load_dict_from_json(file_json):
 
 class CSVWriter:
     """
-    for writing tabular data to a csv file
+    CSVWriter class for writing tabular data to a csv file
+
+    ----------
+    Attributes
+    ----------
+    file_name : str
+        file name of the csv file
+    column_names : list
+        a list of column names
+
     """
     def __init__(self, file_name, column_names):
         self.file_name = file_name
@@ -58,6 +67,12 @@ class CSVWriter:
     def write_row(self, row):
         """
         writes a row into csv file
+
+        ---------
+        Arguments
+        ---------
+        row : list
+            a list of row values
         """
         self.writer.writerow(row)
         return

training/metrics.py

@@ -6,6 +6,21 @@ import torch.nn.functional as F
 
 # compute mean pixel accuracy
 def compute_mean_pixel_acc(true_label, pred_label):
+    """
+    ---------
+    Arguments
+    ---------
+    true_label : ndarray
+        a numpy array of groundtruth label
+    pred_label : ndarray
+        a numpy array of prediction label
+
+    -------
+    Returns
+    -------
+    mean_pixel_accuracy : float
+        mean pixel accuracy
+    """
     if true_label.shape != pred_label.shape:
         print("true_label has dimension", true_label.shape, ", pred_label values have shape", pred_label.shape)
         return
@@ -33,6 +48,23 @@ def compute_mean_pixel_acc(true_label, pred_label):
 
 # compute mean IOU
 def compute_mean_IOU(true_label, pred_label, num_classes=5):
+    """
+    ---------
+    Arguments
+    ---------
+    true_label : ndarray
+        a numpy array of groundtruth label
+    pred_label : ndarray
+        a numpy array of prediction label
+    num_classes : int
+        number of classes in the dataset (default: 5)
+
+    -------
+    Returns
+    -------
+    mean_iou : float
+        mean IoU
+    """
     iou_list = list()
     present_iou_list = list()
 
@@ -55,6 +87,24 @@ def compute_mean_IOU(true_label, pred_label, num_classes=5):
     return np.mean(present_iou_list)
 
 def compute_class_IOU(true_label, pred_label, num_classes=5):
+    """
+    ---------
+    Arguments
+    ---------
+    true_label : ndarray
+        a numpy array of groundtruth label
+    pred_label : ndarray
+        a numpy array of prediction label
+    num_classes : int
+        number of classes in the dataset (default: 5)
+
+
+    -------
+    Returns
+    -------
+    per_class_iou : ndarray
+        a numpy array of per class IoU
+    """
     iou_list = list()
     present_iou_list = list()
 

training/seg_models.py

@@ -8,6 +8,15 @@ from encoder_models import resnet18, resnet34, resnet50, resnet101, efficientnet
 
 class ResNet18DeepLabV3Plus(nn.Module):
     def __init__(self, num_classes, pretrained=True):
+        """
+        ----------
+        Attributes
+        ----------
+        num_classes : int
+            number of classes in the dataset
+        pretrained : bool
+            indicates whether to load pretrained weights for the encoder model (default: True)
+        """
         super().__init__()
 
         self.encoder = resnet18(pretrained=pretrained)
@@ -22,6 +31,15 @@ class ResNet18DeepLabV3Plus(nn.Module):
 
 class ResNet34DeepLabV3Plus(nn.Module):
     def __init__(self, num_classes, pretrained=True):
+        """
+        ----------
+        Attributes
+        ----------
+        num_classes : int
+            number of classes in the dataset
+        pretrained : bool
+            indicates whether to load pretrained weights for the encoder model (default: True)
+        """
         super().__init__()
 
         self.encoder = resnet34(pretrained=pretrained)
@@ -36,6 +54,15 @@ class ResNet34DeepLabV3Plus(nn.Module):
 
 class ResNet50DeepLabV3Plus(nn.Module):
     def __init__(self, num_classes, pretrained=True):
+        """
+        ----------
+        Attributes
+        ----------
+        num_classes : int
+            number of classes in the dataset
+        pretrained : bool
+            indicates whether to load pretrained weights for the encoder model (default: True)
+        """
         super().__init__()
 
         self.encoder = resnet50(pretrained=pretrained)
@@ -50,6 +77,15 @@ class ResNet50DeepLabV3Plus(nn.Module):
 
 class ResNet101DeepLabV3Plus(nn.Module):
     def __init__(self, num_classes, pretrained=True):
+        """
+        ----------
+        Attributes
+        ----------
+        num_classes : int
+            number of classes in the dataset
+        pretrained : bool
+            indicates whether to load pretrained weights for the encoder model (default: True)
+        """
         super().__init__()
 
         self.encoder = resnet101(pretrained=pretrained)
@@ -64,6 +100,15 @@ class ResNet101DeepLabV3Plus(nn.Module):
 
 class EfficientNetSDeepLabV3(nn.Module):
     def __init__(self, num_classes, pretrained=True):
+        """
+        ----------
+        Attributes
+        ----------
+        num_classes : int
+            number of classes in the dataset
+        pretrained : bool
+            indicates whether to load pretrained weights for the encoder model (default: True)
+        """
         super().__init__()
 
         self.encoder = efficientnet_v2_s(pretrained=pretrained)
@@ -81,6 +126,15 @@ class EfficientNetSDeepLabV3(nn.Module):
 
 class EfficientNetMDeepLabV3(nn.Module):
     def __init__(self, num_classes, pretrained=True):
+        """
+        ----------
+        Attributes
+        ----------
+        num_classes : int
+            number of classes in the dataset
+        pretrained : bool
+            indicates whether to load pretrained weights for the encoder model (default: True)
+        """
         super().__init__()
 
         self.encoder = efficientnet_v2_m(pretrained=pretrained)
@@ -98,6 +152,15 @@ class EfficientNetMDeepLabV3(nn.Module):
 
 class EfficientNetLDeepLabV3(nn.Module):
     def __init__(self, num_classes, pretrained=True):
+        """
+        ----------
+        Attributes
+        ----------
+        num_classes : int
+            number of classes in the dataset
+        pretrained : bool
+            indicates whether to load pretrained weights for the encoder model (default: True)
+        """
         super().__init__()
 
         self.encoder = efficientnet_v2_l(pretrained=pretrained)

training/train.py

@@ -19,6 +19,23 @@ from logger_utils import CSVWriter, write_dict_to_json
 from torch.optim.lr_scheduler import _LRScheduler
 
 class PolynomialLR(_LRScheduler):
+    """
+    PolynomialLR class for the polynomial learning rate scheduler
+
+    ----------
+    Attributes
+    ----------
+    optimizer : object
+        object of type optimizer
+    max_epochs : int
+        maximum number of epochs for which optimization needs to be run
+    power : float
+        the power term in the polynomial learning rate scheduler (default: 0.9)
+    last_epoch : int
+        last epoch in the optimization (default: -1)
+    min_lr : float
+        minimum value for the learning rate (default: 1e-6)
+    """
     def __init__(self, optimizer, max_epochs, power=0.9, last_epoch=-1, min_lr=1e-6):
         self.power = power
         self.max_epochs = max_epochs
@@ -30,6 +47,25 @@ class PolynomialLR(_LRScheduler):
                 for base_lr in self.base_lrs]
 
 def validation_loop(dataset_loader, model, ce_loss, device):
+    """
+    ---------
+    Arguments
+    ---------
+    dataset_loader : object
+        object of type dataloader
+    model : object
+        object of type model
+    ce_loss : object
+        object of type cross entropy loss
+    device : str
+        device on which training needs to be run
+
+    -------
+    Returns
+    -------
+    (valid_loss, valid_acc, valid_IOU) : tuple
+        a tuples of torch floats of mean loss, mean accuracy, mean IoU for the validation set
+    """
     model.eval()
     size = len(dataset_loader.dataset)
     num_batches = len(dataset_loader)
@@ -55,6 +91,27 @@ def validation_loop(dataset_loader, model, ce_loss, device):
     return valid_loss, valid_acc, valid_IOU
 
 def train_loop(dataset_loader, model, ce_loss, optimizer, device):
+    """
+    ---------
+    Arguments
+    ---------
+    dataset_loader : object
+        object of type dataloader
+    model : object
+        object of type model
+    ce_loss : object
+        object of type cross entropy loss
+    optimizer : object
+        object of type optimizer
+    device : str
+        device on which training needs to be run
+
+    -------
+    Returns
+    -------
+    train_loss : torch float
+        mean loss for the training set
+    """
     model.train()
     size = len(dataset_loader.dataset)
     num_batches = len(dataset_loader)