Cleanup docstrings in lightly/utils subpackage (#1698)

lightly/utils/bounding_box.py

@@ -1,4 +1,4 @@
-""" Bounding Box Utils """
+"""Bounding Box Utils"""
 
 from __future__ import annotations
 
@@ -31,17 +31,26 @@ class BoundingBox:
         >>> # (x0, y0, x1, y1) = (10, 20, 30, 40)
         >>> W, H = 100, 100 # get image shape
         >>> bbox = BoundingBox(10 / W, 20 / H, 30 / W, 40 / H)
-
     """
 
     def __init__(
         self, x0: float, y0: float, x1: float, y1: float, clip_values: bool = True
     ):
-        """
-        clip_values:
-            Set to true to clip the values into [0, 1] instead of raising an error if they lie outside.
-        """
+        """Initializes a BoundingBox object.
+
+        Args:
+            x0:
+                x0 coordinate relative to image width.
+            y0:
+                y0 coordinate relative to image height.
+            x1:
+                x1 coordinate relative to image width.
+            y1:
+                y1 coordinate relative to image height.
+            clip_values:
+                If True, clips the coordinates to [0, 1].
 
+        """
         if clip_values:
 
             def clip_to_0_1(value: float) -> float:
@@ -60,14 +69,12 @@ def clip_to_0_1(value: float) -> float:
 
         if x0 >= x1:
             raise ValueError(
-                f"x0 must be smaller than x1 for bounding box "
-                f"[{x0}, {y0}, {x1}, {y1}]"
+                f"x0 must be smaller than x1 for bounding box [{x0}, {y0}, {x1}, {y1}]"
             )
 
         if y0 >= y1:
             raise ValueError(
-                "y0 must be smaller than y1 for bounding box "
-                f"[{x0}, {y0}, {x1}, {y1}]"
+                f"y0 must be smaller than y1 for bounding box [{x0}, {y0}, {x1}, {y1}]"
             )
 
         self.x0 = x0
@@ -77,7 +84,20 @@ def clip_to_0_1(value: float) -> float:
 
     @classmethod
     def from_x_y_w_h(cls, x: float, y: float, w: float, h: float) -> BoundingBox:
-        """Helper to convert from bounding box format with width and height.
+        """Creates a BoundingBox from x, y, width, and height.
+
+        Args:
+            x:
+                x coordinate of the top-left corner relative to image width.
+            y:
+                y coordinate of the top-left corner relative to image height.
+            w:
+                Width of the bounding box relative to image width.
+            h:
+                Height of the bounding box relative to image height.
+
+        Returns:
+            BoundingBox: A BoundingBox instance.
 
         Examples:
             >>> bbox = BoundingBox.from_x_y_w_h(0.1, 0.2, 0.2, 0.2)
@@ -89,11 +109,23 @@ def from_x_y_w_h(cls, x: float, y: float, w: float, h: float) -> BoundingBox:
     def from_yolo_label(
         cls, x_center: float, y_center: float, w: float, h: float
     ) -> BoundingBox:
-        """Helper to convert from yolo label format
-        x_center, y_center, w, h --> x0, y0, x1, y1
+        """Creates a BoundingBox from YOLO label format.
+
+        Args:
+            x_center:
+                x coordinate of the center relative to image width.
+            y_center:
+                y coordinate of the center relative to image height.
+            w:
+                Width of the bounding box relative to image width.
+            h:
+                Height of the bounding box relative to image height.
+
+        Returns:
+            BoundingBox: A BoundingBox instance.
 
         Examples:
-            >>> bbox = BoundingBox.from_yolo(0.5, 0.4, 0.2, 0.3)
+            >>> bbox = BoundingBox.from_yolo_label(0.5, 0.4, 0.2, 0.3)
 
         """
         return cls(

lightly/utils/debug.py

@@ -15,7 +15,6 @@
         "'pip install lightly[matplotlib]'."
     )
 except ImportError as ex:
-    # Matplotlib import can fail if an incompatible dateutil version is installed.
     plt = ex
 
 
@@ -24,9 +23,9 @@ def std_of_l2_normalized(z: torch.Tensor) -> torch.Tensor:
     """Calculates the mean of the standard deviation of z along each dimension.
 
     This measure was used by [0] to determine the level of collapse of the
-    learned representations. If the returned number is 0., the outputs z have
-    collapsed to a constant vector. "If the output z has a zero-mean isotropic
-    Gaussian distribution" [0], the returned number should be close to 1/sqrt(d)
+    learned representations. If the returned value is 0., the outputs z have
+    collapsed to a constant vector. If the output z has a zero-mean isotropic
+    Gaussian distribution [0], the returned value should be close to 1/sqrt(d),
     where d is the dimensionality of the output.
 
     [0]: https://arxiv.org/abs/2011.10566
@@ -38,9 +37,7 @@ def std_of_l2_normalized(z: torch.Tensor) -> torch.Tensor:
     Returns:
         The mean of the standard deviation of the l2 normalized tensor z along
         each dimension.
-
     """
-
     if len(z.shape) != 2:
         raise ValueError(
             f"Input tensor must have two dimensions but has {len(z.shape)}!"
@@ -53,8 +50,18 @@ def std_of_l2_normalized(z: torch.Tensor) -> torch.Tensor:
 def apply_transform_without_normalize(
     image: Image.Image,
     transform,
-):
-    """Applies the transform to the image but skips ToTensor and Normalize."""
+) -> Image.Image:
+    """Applies the transform to the image but skips ToTensor and Normalize.
+
+    Args:
+        image:
+            The input PIL image.
+        transform:
+            The transformation to apply, excluding ToTensor and Normalize.
+
+    Returns:
+        The transformed image.
+    """
     skippable_transforms = (
         torchvision.transforms.ToTensor,
         torchvision.transforms.Normalize,
@@ -70,10 +77,10 @@ def apply_transform_without_normalize(
 def generate_grid_of_augmented_images(
     input_images: List[Image.Image],
     collate_function: Union[BaseCollateFunction, MultiViewCollateFunction],
-):
+) -> List[List[Image.Image]]:
     """Returns a grid of augmented images. Images in a column belong together.
 
-    This function ignores the transforms ToTensor and Normalize for visualization purposes.
+    This function ignores the ToTensor and Normalize transforms for visualization purposes.
 
     Args:
         input_images:
@@ -116,9 +123,9 @@ def plot_augmented_images(
     input_images: List[Image.Image],
     collate_function: Union[BaseCollateFunction, MultiViewCollateFunction],
 ):
-    """Returns a figure showing original images in the left column and augmented images to their right.
+    """Plots original images and augmented images in a figure.
 
-    This function ignores the transforms ToTensor and Normalize for visualization purposes.
+    This function ignores the ToTensor and Normalize transforms for visualization purposes.
 
     Args:
         input_images:
@@ -134,7 +141,6 @@ def plot_augmented_images(
         MultiViewCollateFunctions all the generated views are shown.
 
     """
-
     _check_matplotlib_available()
 
     if len(input_images) == 0:
@@ -166,5 +172,6 @@ def plot_augmented_images(
 
 
 def _check_matplotlib_available() -> None:
+    """Checks if matplotlib is available. Raises an error if not."""
     if isinstance(plt, Exception):
         raise plt

lightly/utils/dependency.py

@@ -3,24 +3,42 @@
 
 @functools.lru_cache(maxsize=1)
 def torchvision_vit_available() -> bool:
+    """Checks if Vision Transformer (ViT) models are available in torchvision.
+
+    This function checks if the `vision_transformer` module is available in torchvision,
+    which requires torchvision version >= 0.12. It also handles exceptions related to
+    CUDA version mismatches and installation issues.
+
+    Returns:
+        True if the Vision Transformer (ViT) models are available in torchvision,
+        otherwise False.
+    """
     try:
-        import torchvision.models.vision_transformer  # Requires torchvision >=0.12
+        import torchvision.models.vision_transformer  # Requires torchvision >=0.12.
     except (
-        RuntimeError,  # Different CUDA versions for torch and torchvision
-        OSError,  # Different CUDA versions for torch and torchvision (old)
-        ImportError,  # No installation or old version of torchvision
+        RuntimeError,  # Different CUDA versions for torch and torchvision.
+        OSError,  # Different CUDA versions for torch and torchvision (old).
+        ImportError,  # No installation or old version of torchvision.
     ):
         return False
-    else:
-        return True
+    return True
 
 
 @functools.lru_cache(maxsize=1)
 def timm_vit_available() -> bool:
+    """Checks if Vision Transformer (ViT) models are available in the timm library.
+
+    This function checks if the `vision_transformer` module and `LayerType` from timm
+    are available, which requires timm version >= 0.3.3 and >= 0.9.9, respectively.
+
+    Returns:
+        True if the Vision Transformer (ViT) models are available in timm,
+        otherwise False.
+
+    """
     try:
         import timm.models.vision_transformer  # Requires timm >= 0.3.3
         from timm.layers import LayerType  # Requires timm >= 0.9.9
     except ImportError:
         return False
-    else:
-        return True
+    return True

lightly/utils/dist.py

@@ -8,19 +8,19 @@
 class GatherLayer(torch.autograd.Function):
     """Gather tensors from all processes, supporting backward propagation.
 
-    This code was taken and adapted from here:
+    Adapted from the Solo-Learn project:
     https://github.com/vturrisi/solo-learn/blob/b69b4bd27472593919956d9ac58902a301537a4d/solo/utils/misc.py#L187
 
     """
 
     @staticmethod
-    def forward(ctx, input: torch.Tensor) -> Tuple[torch.Tensor, ...]:  # type: ignore
+    def forward(ctx: FunctionCtx, input: torch.Tensor) -> Tuple[torch.Tensor, ...]:  # type: ignore
         output = [torch.empty_like(input) for _ in range(dist.get_world_size())]
         dist.all_gather(output, input)
         return tuple(output)
 
     @staticmethod
-    def backward(ctx, *grads) -> torch.Tensor:  # type: ignore
+    def backward(ctx: FunctionCtx, *grads: torch.Tensor) -> torch.Tensor:  # type: ignore
         all_gradients = torch.stack(grads)
         dist.all_reduce(all_gradients)
         grad_out = all_gradients[dist.get_rank()]
@@ -38,7 +38,7 @@ def world_size() -> int:
 
 
 def gather(input: torch.Tensor) -> Tuple[torch.Tensor]:
-    """Gathers this tensor from all processes. Supports backprop."""
+    """Gathers a tensor from all processes and supports backpropagation."""
     return GatherLayer.apply(input)  # type: ignore[no-any-return]
 
 
@@ -62,6 +62,9 @@ def eye_rank(n: int, device: Optional[torch.device] = None) -> torch.Tensor:
         device:
             Device on which the matrix should be created.
 
+    Returns:
+        A tensor with the appropriate diagonal filled for this rank.
+
     """
     rows = torch.arange(n, device=device, dtype=torch.long)
     cols = rows + rank() * n
@@ -74,7 +77,7 @@ def eye_rank(n: int, device: Optional[torch.device] = None) -> torch.Tensor:
 
 
 def rank_zero_only(fn: Callable[..., R]) -> Callable[..., Optional[R]]:
-    """Decorator that only runs the function on the process with rank 0.
+    """Decorator to ensure the function only runs on the process with rank 0.
 
     Example:
         >>> @rank_zero_only

lightly/utils/embeddings_2d.py

@@ -1,4 +1,4 @@
-""" Transform embeddings to two-dimensional space for visualization. """
+"""Transforms embeddings to two-dimensional space for visualization."""
 
 # Copyright (c) 2020. Lightly AG and its affiliates.
 # All Rights Reserved
@@ -21,13 +21,18 @@ class PCA(object):
             Number of principal components to keep.
         eps:
             Epsilon for numerical stability.
+        mean:
+            Mean of the data.
+        w:
+            Eigenvectors of the covariance matrix.
+
     """
 
     def __init__(self, n_components: int = 2, eps: float = 1e-10):
         self.n_components = n_components
+        self.eps = eps
         self.mean: Optional[NDArray[np.float32]] = None
         self.w: Optional[NDArray[np.float32]] = None
-        self.eps = eps
 
     def fit(self, X: NDArray[np.float32]) -> PCA:
         """Fits PCA to data in X.
@@ -37,7 +42,7 @@ def fit(self, X: NDArray[np.float32]) -> PCA:
                 Datapoints stored in numpy array of size n x d.
 
         Returns:
-            PCA object to transform datapoints.
+            PCA: The fitted PCA object to transform data points.
 
         """
         X = X.astype(np.float32)
@@ -46,7 +51,7 @@ def fit(self, X: NDArray[np.float32]) -> PCA:
         X = X - self.mean + self.eps
         cov = np.cov(X.T) / X.shape[0]
         v, w = np.linalg.eig(cov)
-        idx = v.argsort()[::-1]
+        idx = v.argsort()[::-1]  # Sort eigenvalues in descending order
         v, w = v[idx], w[:, idx]
         self.w = w
         return self
@@ -62,10 +67,13 @@ def transform(self, X: NDArray[np.float32]) -> NDArray[np.float32]:
             Numpy array of n x p datapoints where p <= d.
 
         Raises:
-            ValueError: If PCA was not fitted before.
+            ValueError:
+                If PCA is not fitted before calling this method.
+
         """
         if self.mean is None or self.w is None:
             raise ValueError("PCA not fitted yet. Call fit() before transform().")
+
         X = X.astype(np.float32)
         X = X - self.mean + self.eps
         transformed: NDArray[np.float32] = X.dot(self.w)[:, : self.n_components]
@@ -77,7 +85,7 @@ def fit_pca(
     n_components: int = 2,
     fraction: Optional[float] = None,
 ) -> PCA:
-    """Fits PCA to randomly selected subset of embeddings.
+    """Fits PCA to a randomly selected subset of embeddings.
 
     For large datasets, it can be unfeasible to perform PCA on the whole data.
     This method can fit a PCA on a fraction of the embeddings in order to save
@@ -101,8 +109,7 @@ def fit_pca(
     """
     if fraction is not None:
         if fraction < 0.0 or fraction > 1.0:
-            msg = f"fraction must be in [0, 1] but was {fraction}."
-            raise ValueError(msg)
+            raise ValueError(f"fraction must be in [0, 1] but was {fraction}.")
 
     N = embeddings.shape[0]
     n = N if fraction is None else min(N, int(N * fraction))