ADD: Add docstrings

Author: adamshephard

tiatoolbox/models/architecture/vanilla.py

@@ -22,23 +22,32 @@ def _get_architecture(
     weights: str or WeightsEnum = "DEFAULT",
     **kwargs: dict,
 ) -> list[nn.Sequential, ...] | nn.Sequential:
-    """Get a model.
+    """Retrieve a CNN model architecture.
 
-    Model architectures are either already defined within torchvision or
-    they can be custom-made within tiatoolbox.
+    This function fetches a Convolutional Neural Network (CNN) model architecture,
+    either predefined in torchvision or custom-made within tiatoolbox, for
+    patch classification tasks.
 
     Args:
         arch_name (str):
-            Architecture name.
+            Name of the architecture (e.g. 'resnet50', 'alexnet').
         weights (str or WeightsEnum):
-            torchvision model weights (get_model_weights).
-        kwargs (dict):
+            Pretrained torchvision model weights to use (get_model_weights).
+            Defaults to "DEFAULT".
+        **kwargs (dict):
             Key-word arguments.
 
     Returns:
-        List of PyTorch network layers wrapped with `nn.Sequential`.
-        https://pytorch.org/docs/stable/generated/torch.nn.Sequential.html
+        list[nn.Sequential, ...] | nn.Sequential:
+            A list of PyTorch network layers wrapped with `nn.Sequential`.
 
+    Raises:
+        ValueError:
+            If `arch_name` is not supported.
+
+    Example:
+        >>> model = _get_architecture("resnet18")
+        >>> print(model)
     """
     backbone_dict = {
         "alexnet": torch_models.alexnet,
@@ -85,7 +94,10 @@ def _get_timm_architecture(
     *,
     pretrained: bool,
 ) -> list[nn.Sequential, ...] | nn.Sequential:
-    """Get architecture and weights for pathology-specific timm models.
+    """Retrieve a timm model architecture.
+
+    This function fetches a model architecture from the timm library, specifically for
+    pathology-related tasks.
 
     Args:
         arch_name (str):
@@ -94,12 +106,16 @@ def _get_timm_architecture(
             Whether to load pretrained weights.
 
     Returns:
-        A ready-to-use timm model.
+        list[nn.Sequential, ...] | nn.Sequential:
+            A ready-to-use timm model.
 
     Raises:
         ValueError:
             If the backbone architecture is not supported.
 
+    Example:
+        >>> model = _get_timm_architecture("UNI", pretrained=True)
+        >>> print(model)
     """
     if arch_name in [f"efficientnet_b{i}" for i in range(8)]:
         model = timm.create_model(arch_name, pretrained=pretrained)
@@ -177,6 +193,13 @@ def _postproc(image: np.ndarray) -> np.ndarray:
 
     This simply applies argmax along last axis of the input.
 
+    Args:
+        image (np.ndarray):
+            The input image array.
+
+    Returns:
+        np.ndarray:
+            The post-processed image array.
     """
     return np.argmax(image, axis=-1)
 
@@ -197,8 +220,15 @@ def _infer_batch(
             A batch of data generated by
             `torch.utils.data.DataLoader`.
         device (str):
-                Transfers model to the specified device. Default is "cpu".
+            Transfers model to the specified device. Default is "cpu".
+
+    Returns:
+        dict[str, np.ndarray]:
+            The model predictions as a NumPy array.
 
+    Example:
+        >>> output = _infer_batch(model, batch_data, "cuda")
+        >>> print(output)
     """
     img_patches_device = batch_data.to(device=device).type(
         torch.float32,
@@ -217,11 +247,14 @@ def _infer_batch(
 class CNNModel(ModelABC):
     """Retrieve the model backbone and attach an extra FCN to perform classification.
 
+    This class initializes a Convolutional Neural Network (CNN) model with a specified
+    backbone and attaches a fully connected layer for classification tasks.
+
     Args:
         backbone (str):
-            Model name.
+            Name of the CNN model backbone (e.g. "resnet18", "densenet121").
         num_classes (int):
-            Number of classes output by model.
+            Number of classes output by model. Defaults to 1.
 
     Attributes:
         num_classes (int):
@@ -231,9 +264,12 @@ class CNNModel(ModelABC):
         pool (nn.Module):
             Type of pooling applied after feature extraction.
         classifier (nn.Module):
-            Linear classifier module used to map the features to the
-            output.
+            Linear classifier module used to map the features to the output.
 
+    Example:
+        >>> model = CNNModel("resnet18", num_classes=2)
+        >>> output = model(torch.randn(1, 3, 224, 224))
+        >>> print(output.shape)
     """
 
     def __init__(self: CNNModel, backbone: str, num_classes: int = 1) -> None:
@@ -257,6 +293,9 @@ def forward(self: CNNModel, imgs: torch.Tensor) -> torch.Tensor:
             imgs (torch.Tensor):
                 Model input.
 
+        Returns:
+            torch.Tensor:
+                The output logits after passing through the model.
         """
         feat = self.feat_extract(imgs)
         gap_feat = self.pool(feat)
@@ -270,6 +309,13 @@ def postproc(image: np.ndarray) -> np.ndarray:
 
         This simply applies argmax along last axis of the input.
 
+        Args:
+            image (np.ndarray):
+                The input image array.
+
+        Returns:
+            np.ndarray:
+                The post-processed image array.
         """
         return _postproc(image=image)
 
@@ -292,6 +338,9 @@ def infer_batch(
             device (str):
                 Transfers model to the specified device. Default is "cpu".
 
+        Example:
+            >>> output = _infer_batch(model, batch_data, "cuda")
+            >>> print(output)
         """
         return _infer_batch(model=model, batch_data=batch_data, device=device)
 
@@ -327,6 +376,11 @@ class TimmModel(ModelABC):
         classifier (nn.Module):
             Linear classifier module used to map the features to the
             output.
+
+    Example:
+        >>> model = TimmModel("UNI", pretrained=True)
+        >>> output = model(torch.randn(1, 3, 224, 224))
+        >>> print(output.shape)
     """
 
     def __init__(
@@ -357,6 +411,9 @@ def forward(self: TimmModel, imgs: torch.Tensor) -> torch.Tensor:
             imgs (torch.Tensor):
                 Model input.
 
+        Returns:
+            torch.Tensor:
+                The output logits after passing through the model.
         """
         feat = self.feat_extract(imgs)
         feat = torch.flatten(feat, 1)
@@ -369,6 +426,14 @@ def postproc(image: np.ndarray) -> np.ndarray:
 
         This simply applies argmax along last axis of the input.
 
+        Args:
+            image (np.ndarray):
+                The input image array.
+
+        Returns:
+            np.ndarray:
+                The post-processed image array.
+
         """
         return _postproc(image=image)
 
@@ -391,6 +456,14 @@ def infer_batch(
             device (str):
                 Transfers model to the specified device. Default is "cpu".
 
+        Returns:
+            dict[str, np.ndarray]:
+                The model predictions as a NumPy array.
+
+        Example:
+            >>> output = _infer_batch(model, batch_data, "cuda")
+            >>> print(output)
+
         """
         return _infer_batch(model=model, batch_data=batch_data, device=device)
 
@@ -423,6 +496,12 @@ class CNNBackbone(ModelABC):
                 - "mobilenet_v3_large"
                 - "mobilenet_v3_small"
 
+    Attributes:
+        feat_extract (nn.Module):
+            Backbone CNN model.
+        pool (nn.Module):
+            Type of pooling applied after feature extraction.
+
     Examples:
         >>> # Creating resnet50 architecture from default pytorch
         >>> # without the classification layer with its associated
@@ -452,6 +531,9 @@ def forward(self: CNNBackbone, imgs: torch.Tensor) -> torch.Tensor:
             imgs (torch.Tensor):
                 Model input.
 
+        Returns:
+            torch.Tensor:
+                The extracted features.
         """
         feat = self.feat_extract(imgs)
         gap_feat = self.pool(feat)
@@ -480,6 +562,9 @@ def infer_batch(
             list[dict[str, np.ndarray]]:
                 list of dictionary values with numpy arrays.
 
+        Example:
+            >>> output = CNNBackbone.infer_batch(model, batch_data, "cuda")
+            >>> print(output)
         """
         return [_infer_batch(model=model, batch_data=batch_data, device=device)]
 
@@ -491,8 +576,7 @@ class TimmBackbone(ModelABC):
 
     Args:
         backbone (str):
-            Model name. Currently, the tool supports following
-              model names and their default associated weights from timm.
+            Model name. Supported model names include:
                 - "efficientnet_b{i}" for i in [0, 1, ..., 7]
                 - "UNI"
                 - "prov-gigapath"
@@ -503,6 +587,10 @@ class TimmBackbone(ModelABC):
         pretrained (bool, keyword-only):
             Whether to load pretrained weights.
 
+    Attributes:
+        feat_extract (nn.Module):
+            Backbone timm model.
+
     Examples:
         >>> # Creating UNI tile encoder
         >>> model = TimmBackbone(backbone="UNI", pretrained=True)
@@ -531,6 +619,9 @@ def forward(self: TimmBackbone, imgs: torch.Tensor) -> torch.Tensor:
             imgs (torch.Tensor):
                 Model input.
 
+        Returns:
+            torch.Tensor:
+                The extracted features.
         """
         feats = self.feat_extract(imgs)
         return torch.flatten(feats, 1)
@@ -558,5 +649,8 @@ def infer_batch(
             list[dict[str, np.ndarray]]:
                 list of dictionary values with numpy arrays.
 
+        Example:
+            >>> output = TimmBackbone.infer_batch(model, batch_data, "cuda")
+            >>> print(output)
         """
         return [_infer_batch(model=model, batch_data=batch_data, device=device)]