Fix docstrings

anomalib/config/config.py

@@ -35,7 +35,7 @@ def update_input_size_config(config: Union[DictConfig, ListConfig]) -> Union[Dic
         config (Union[DictConfig, ListConfig]): Configurable parameters object
 
     Returns:
-        Configurable parameters with updated values
+        Union[DictConfig, ListConfig]: Configurable parameters with updated values
     """
     # handle image size
     if isinstance(config.dataset.image_size, int):
@@ -56,10 +56,10 @@ def update_nncf_config(config: Union[DictConfig, ListConfig]) -> Union[DictConfi
     """Set the NNCF input size based on the value of the crop_size parameter in the configurable parameters object.
 
     Args:
-        config: Dictconfig: Configurable parameters of the current run.
+        config (Union[DictConfig, ListConfig]): Configurable parameters of the current run.
 
     Returns:
-        Updated configurable parameters in DictConfig object.
+        Union[DictConfig, ListConfig]: Updated configurable parameters in DictConfig object.
     """
     crop_size = config.dataset.image_size
     sample_size = (crop_size, crop_size) if isinstance(crop_size, int) else crop_size
@@ -157,7 +157,7 @@ def get_configurable_parameters(
         config_file_extension: Optional[str]:  (Default value = "yaml")
 
     Returns:
-        Configurable parameters in DictConfig object.
+        Union[DictConfig, ListConfig]: Configurable parameters in DictConfig object.
     """
     if model_name is None and model_config_path is None:
         raise ValueError(

anomalib/core/callbacks/compress.py

@@ -9,7 +9,7 @@
 class CompressModelCallback(Callback):
     """Callback to compresses a trained model.
 
-    Model is first exported to .onnx format, and then converted to OpenVINO IR.
+    Model is first exported to ``.onnx`` format, and then converted to OpenVINO IR.
 
     Args:
         input_size (Tuple[int, int]): Tuple of image height, width
@@ -23,7 +23,11 @@ def __init__(self, input_size: Tuple[int, int], dirpath: str, filename: str):
         self.filename = filename
 
     def on_train_end(self, trainer, pl_module: LightningModule) -> None:  # pylint: disable=W0613
-        """Call when the train ends."""
+        """Call when the train ends.
+
+        Converts the model to ``onnx`` format and then calls OpenVINO's model optimizer to get the
+        ``.xml`` and ``.bin`` IR files.
+        """
         os.makedirs(self.dirpath, exist_ok=True)
         onnx_path = os.path.join(self.dirpath, self.filename + ".onnx")
         height, width = self.input_size

anomalib/core/callbacks/model_loader.py

@@ -10,5 +10,8 @@ def __init__(self, weights_path):
         self.weights_path = weights_path
 
     def on_test_start(self, trainer, pl_module: LightningModule) -> None:  # pylint: disable=W0613
-        """Call when the test begins."""
+        """Call when the test begins.
+
+        Loads the model weights from into the PyTorch module.
+        """
         pl_module.load_state_dict(torch.load(self.weights_path)["state_dict"])

anomalib/core/callbacks/nncf_callback.py

@@ -1,15 +1,15 @@
 """NNCF Callback."""
 
 import os
-from typing import Any, Dict, Iterator, Optional, Tuple
+from typing import Any, Dict, Iterator, Optional, Tuple, Union
 
 import pytorch_lightning as pl
 import yaml
 from nncf import NNCFConfig
 from nncf.api.compression import CompressionAlgorithmController, CompressionScheduler
 from nncf.torch import create_compressed_model, register_default_init_args
 from nncf.torch.initialization import PTInitializingDataLoader
-from omegaconf import OmegaConf
+from omegaconf import DictConfig, ListConfig, OmegaConf
 from pytorch_lightning import Callback
 from torch.utils.data.dataloader import DataLoader
 
@@ -50,6 +50,8 @@ def get_inputs(self, dataloader_output) -> Tuple[Tuple, Dict]:
     def get_target(self, _):
         """Return structure for ground truth in loss criterion based on dataloader output.
 
+        This implementation does not do anything and is a placeholder.
+
         Returns:
             None
         """
@@ -61,9 +63,14 @@ class NNCFCallback(Callback):
 
     Assumes that the pl module contains a 'model' attribute, which is
     the PyTorch module that must be compressed.
+
+    Args:
+        config (Union[ListConfig, DictConfig]): NNCF Configuration
+        dirpath (str): Path where the export `onnx` and the OpenVINO `xml` and `bin` IR are saved.
+        filename (str): Name of the generated model files.
     """
 
-    def __init__(self, config, dirpath, filename):
+    def __init__(self, config: Union[ListConfig, DictConfig], dirpath: str, filename: str):
         config_dict = yaml.safe_load(OmegaConf.to_yaml(config.optimization.nncf))
         self.nncf_config = NNCFConfig.from_dict(config_dict)
         self.dirpath = dirpath
@@ -78,7 +85,10 @@ def __init__(self, config, dirpath, filename):
         self.compression_scheduler: CompressionScheduler
 
     def setup(self, _: pl.Trainer, pl_module: pl.LightningModule, __: Optional[str] = None) -> None:
-        """Call when fit or test begins."""
+        """Call when fit or test begins.
+
+        Takes the pytorch model and wraps it using the compression controller so that it is ready for nncf fine-tuning.
+        """
         if self.comp_ctrl is None:
             init_loader = InitLoader(self.train_loader)
             nncf_config = register_default_init_args(
@@ -91,22 +101,29 @@ def setup(self, _: pl.Trainer, pl_module: pl.LightningModule, __: Optional[str]
     def on_train_batch_start(
         self, trainer, _pl_module: pl.LightningModule, _batch: Any, _batch_idx: int, _dataloader_idx: int
     ) -> None:
-        """Call when the train batch begins."""
+        """Call when the train batch begins.
+
+        Prepare compression method to continue training the model in the next step.
+        """
         self.compression_scheduler.step()
         if self.comp_ctrl is not None:
             trainer.model.loss_val = self.comp_ctrl.loss()
 
     def on_train_end(self, _trainer, _pl_module: pl.LightningModule) -> None:
-        """Call when the train ends."""
+        """Call when the train ends.
+
+        Exports onnx model and if compression controller is not None, uses the onnx model to generate the OpenVINO IR.
+        """
         os.makedirs(self.dirpath, exist_ok=True)
         onnx_path = os.path.join(self.dirpath, self.filename + ".onnx")
         if self.comp_ctrl is not None:
             self.comp_ctrl.export_model(onnx_path)
         optimize_command = "mo --input_model " + onnx_path + " --output_dir " + self.dirpath
         os.system(optimize_command)
 
-    def on_train_epoch_end(
-        self, _trainer: pl.Trainer, _pl_module: pl.LightningModule, _unused: Optional[Any] = None
-    ) -> None:
-        """Call when the train epoch ends."""
+    def on_train_epoch_start(self, _trainer: pl.Trainer, _pl_module: pl.LightningModule) -> None:
+        """Call when the train epoch starts.
+
+        Prepare compression method to continue training the model in the next epoch.
+        """
         self.compression_scheduler.epoch_step()

anomalib/core/callbacks/save_to_csv.py

@@ -11,15 +11,14 @@
 class SaveToCSVCallback(Callback):
     """Callback that saves the inference results of a model.
 
-    The callback generates a csv file that saves different performance
-    metrics and results.
+    The callback generates a csv file that saves the predicted label, the true label and the image name.
     """
 
     def __init__(self):
         """Callback to save metrics to CSV."""
 
     def on_test_epoch_end(self, _trainer: Trainer, pl_module: AnomalyModule) -> None:
-        """Save Results at the end of training.
+        """Save Results at the end of testing.
 
         Args:
             _trainer (Trainer): Pytorch lightning trainer object (unused)

anomalib/core/callbacks/timer.py

@@ -7,6 +7,7 @@
 class TimerCallback(Callback):
     """Callback that measures the training and testing time of a PyTorch Lightning module."""
 
+    # pylint: disable=unused-argument
     def __init__(self):
         self.start: float
         self.num_images: int = 0

anomalib/core/callbacks/visualizer_callback.py

@@ -21,7 +21,7 @@ class VisualizerCallback(Callback):
     The callback generates a figure showing the original image, the ground truth segmentation mask,
     the predicted error heat map, and the predicted segmentation mask.
 
-    To save the images to the filesystem, add the 'local' keyword to the project.log_images_to parameter in the
+    To save the images to the filesystem, add the 'local' keyword to the `project.log_images_to` parameter in the
     config.yaml file.
     """
 
@@ -34,6 +34,16 @@ def _add_images(
         module: AnomalyModule,
         filename: Path,
     ):
+        """Save image to logger/local storage.
+
+        Saves the image in `visualizer.figure` to the respective loggers and local storage if specified in
+        `log_images_to` in `config.yaml` of the models.
+
+        Args:
+            visualizer (Visualizer): Visualizer object from which the `figure` is saved/logged.
+            module (AnomalyModule): Anomaly module which holds reference to `hparams` and `logger`.
+            filename (Path): Path of the input image. This name is used as name for the generated image.
+        """
 
         # store current logger type as a string
         logger_type = type(module.logger).__name__.lower()

anomalib/core/model/anomaly_module.py

@@ -20,7 +20,7 @@
 import torch
 from omegaconf import DictConfig, ListConfig
 from pytorch_lightning.callbacks.base import Callback
-from torch import nn
+from torch import Tensor, nn
 from torchmetrics import F1, MetricCollection
 
 from anomalib.core.metrics import AUROC, OptimalF1
@@ -29,6 +29,8 @@
 class AnomalyModule(pl.LightningModule):
     """AnomalyModule to train, validate, predict and test images.
 
+    Acts as a base class for all the Anomaly Modules in the library.
+
     Args:
         params (Union[DictConfig, ListConfig]): Configuration
     """
@@ -39,10 +41,10 @@ def __init__(self, params: Union[DictConfig, ListConfig]):
         # Force the type for hparams so that it works with OmegaConfig style of accessing
         self.hparams: Union[DictConfig, ListConfig]  # type: ignore
         self.save_hyperparameters(params)
-        self.loss: torch.Tensor
+        self.loss: Tensor
         self.callbacks: List[Callback]
         self.register_buffer("threshold", torch.tensor(params.model.threshold.default))  # pylint: disable=not-callable
-        self.threshold: torch.Tensor
+        self.threshold: Tensor
 
         self.model: nn.Module
 
@@ -64,20 +66,20 @@ def forward(self, batch):  # pylint: disable=arguments-differ
             batch (Tensor): Input Tensor
 
         Returns:
-            [Tensor]: Output tensor from the model.
+            Tensor: Output tensor from the model.
         """
         return self.model(batch)
 
-    def predict_step(self, batch, batch_idx, _):  # pylint: disable=arguments-differ, signature-differs
+    def predict_step(self, batch, batch_idx, dataloader_idx):  # pylint: disable=arguments-differ, signature-differs
         """Step function called during :meth:`~pytorch_lightning.trainer.trainer.Trainer.predict`.
 
         By default, it calls :meth:`~pytorch_lightning.core.lightning.LightningModule.forward`.
         Override to add any processing logic.
 
         Args:
-            batch: Current batch
-            batch_idx: Index of current batch
-            dataloader_idx: Index of the current dataloader
+            batch (Tensor): Current batch
+            batch_idx (int): Index of current batch
+            dataloader_idx (int): Index of the current dataloader
 
         Return:
             Predicted output
@@ -88,7 +90,7 @@ def test_step(self, batch, _):  # pylint: disable=arguments-differ
         """Calls validation_step for anomaly map/score calculation.
 
         Args:
-          batch: Input batch
+          batch (Tensor): Input batch
           _: Index of the batch.
 
         Returns:
@@ -132,7 +134,7 @@ def _post_process(self, outputs, predict_labels=False):
         """Compute labels based on model predictions."""
         if "pred_scores" not in outputs and "anomaly_maps" in outputs:
             outputs["pred_scores"] = (
-                outputs["anomaly_maps"].reshape(outputs["anomaly_maps"].shape[0], -1).max(axis=1).values
+                outputs["anomaly_maps"].reshape(outputs["anomaly_maps"].shape[0], -1).max(dim=1).values
             )
         if predict_labels:
             outputs["pred_labels"] = outputs["pred_scores"] >= self.threshold.item()

anomalib/core/model/feature_extractor.py

@@ -26,6 +26,10 @@
 class FeatureExtractor(nn.Module):
     """Extract features from a CNN.
 
+    Args:
+        backbone (nn.Module): The backbone to which the feature extraction hooks are attached.
+        layers (Iterable[str]): List of layer names of the backbone to which the hooks are attached.
+
     Example:
         >>> import torch
         >>> import torchvision
@@ -55,7 +59,7 @@ def get_features(self, layer_id: str) -> Callable:
         """Get layer features.
 
         Args:
-            layer_id: str: Layer ID
+            layer_id (str): Layer ID
 
         Returns:
             Layer features

anomalib/core/model/kde.py

@@ -18,32 +18,37 @@
 from typing import Optional
 
 import torch
+from torch import Tensor
 
 from anomalib.core.model.dynamic_module import DynamicBufferModule
 
 
 class GaussianKDE(DynamicBufferModule):
-    """Gaussian Kernel Density Estimation."""
+    """Gaussian Kernel Density Estimation.
 
-    def __init__(self, dataset: Optional[torch.Tensor] = None):
+    Args:
+        dataset (Optional[Tensor], optional): Dataset on which to fit the KDE model. Defaults to None.
+    """
+
+    def __init__(self, dataset: Optional[Tensor] = None):
         super().__init__()
 
         if dataset is not None:
             self.fit(dataset)
 
-        self.register_buffer("bw_transform", torch.Tensor())
-        self.register_buffer("dataset", torch.Tensor())
-        self.register_buffer("norm", torch.Tensor())
+        self.register_buffer("bw_transform", Tensor())
+        self.register_buffer("dataset", Tensor())
+        self.register_buffer("norm", Tensor())
 
-        self.bw_transform = torch.Tensor()
-        self.dataset = torch.Tensor()
-        self.norm = torch.Tensor()
+        self.bw_transform = Tensor()
+        self.dataset = Tensor()
+        self.norm = Tensor()
 
-    def forward(self, features: torch.Tensor) -> torch.Tensor:
+    def forward(self, features: Tensor) -> Tensor:
         """Get the KDE estimates from the feature map.
 
         Args:
-          features: torch.Tensor: Feature map extracted from the CNN
+          features (Tensor): Feature map extracted from the CNN
 
         Returns: KDE Estimates
         """
@@ -57,11 +62,11 @@ def forward(self, features: torch.Tensor) -> torch.Tensor:
 
         return estimate
 
-    def fit(self, dataset: torch.Tensor) -> None:
+    def fit(self, dataset: Tensor) -> None:
         """Fit a KDE model to the input dataset.
 
         Args:
-          dataset: torch.Tensor: Input dataset.
+          dataset (Tensor): Input dataset.
 
         Returns:
             None
@@ -88,12 +93,12 @@ def fit(self, dataset: torch.Tensor) -> None:
         self.norm = norm
 
     @staticmethod
-    def cov(tensor: torch.Tensor, bias: Optional[bool] = False) -> torch.Tensor:
+    def cov(tensor: Tensor, bias: Optional[bool] = False) -> Tensor:
         """Calculate covariance matrix.
 
         Args:
-            tensor: torch.Tensor: Input tensor from which covariance matrix is computed.
-            bias: Optional[bool]:  (Default value = False)
+            tensor (Tensor): Input tensor from which covariance matrix is computed.
+            bias (Optional[bool]):  (Default value = False)
 
         Returns:
             Output covariance matrix.