Speedup in fpixel.grayscale_to_multichannel (#2564)

Author: ternaus

.pre-commit-config.yaml

@@ -70,7 +70,7 @@ repos:
         language: system
         files: setup.py
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.11.13
+    rev: v0.12.0
     hooks:
       - id: ruff
         exclude: '__pycache__/'
@@ -98,7 +98,7 @@ repos:
     hooks:
       - id: pyproject-fmt
   - repo: https://github.com/pre-commit/mirrors-mypy
-    rev: v1.16.0
+    rev: v1.16.1
     hooks:
       - id: mypy
         files: ^albumentations/

albumentations/augmentations/geometric/transforms.py

@@ -1465,7 +1465,7 @@ def get_params_dependent_on_data(
                 for tile in tiles
             ],
         ).reshape(
-            self.num_grid_xy[::-1] + (4,),
+            (*self.num_grid_xy[::-1], 4),
         )  # Reshape to (grid_height, grid_width, 4)
 
         polygons = fgeometric.generate_distorted_grid_polygons(

albumentations/augmentations/pixel/functional.py

@@ -1647,8 +1647,8 @@ def grayscale_to_multichannel(
         return grayscale_image
 
     squeezed = np.squeeze(grayscale_image)
-    # For multi-channel output, stack channels
-    return np.stack([squeezed] * num_output_channels, axis=-1)
+    # For multi-channel output, use tile for better performance
+    return np.tile(squeezed[..., np.newaxis], (1,) * squeezed.ndim + (num_output_channels,))
 
 
 @preserve_channel_dim
@@ -2519,7 +2519,7 @@ def generate_noise(
     height, width = shape[:2]
     reduced_height = max(1, int(height * approximation))
     reduced_width = max(1, int(width * approximation))
-    reduced_shape = (reduced_height, reduced_width) + shape[2:]
+    reduced_shape = (reduced_height, reduced_width, *shape[2:])
 
     # Generate noise at reduced resolution
     if spatial_mode == "shared":
@@ -3482,7 +3482,7 @@ def prepare_drop_values(
         return np.full(array.shape, values[0], dtype=array.dtype)
 
     # For multichannel input, broadcast values to full shape
-    return np.full(array.shape[:2] + (len(values),), values, dtype=array.dtype)
+    return np.full((*array.shape[:2], len(values)), values, dtype=array.dtype)
 
 
 def get_mask_array(data: dict[str, Any]) -> np.ndarray | None:

albumentations/augmentations/pixel/transforms.py

@@ -3678,14 +3678,53 @@ def apply(self, img: np.ndarray, **params: Any) -> np.ndarray:
             warnings.warn("The image is already an RGB.", stacklevel=2)
             return np.ascontiguousarray(img)
         if not is_grayscale_image(img):
-            msg = "ToRGB transformation expects 2-dim images or 3-dim with the last dimension equal to 1."
+            msg = "ToRGB transformation expects images with the number of channels equal to 1."
             raise TypeError(msg)
 
         return fpixel.grayscale_to_multichannel(
             img,
             num_output_channels=self.num_output_channels,
         )
 
+    def apply_to_images(self, images: np.ndarray, **params: Any) -> np.ndarray:
+        """Apply ToRGB to a batch of images.
+
+        Args:
+            images (np.ndarray): Batch of images with shape (N, H, W, C) or (N, H, W).
+            **params (Any): Additional parameters.
+
+        Returns:
+            np.ndarray: Batch of RGB images.
+
+        """
+        return self.apply(images, **params)
+
+    def apply_to_volume(self, volume: np.ndarray, **params: Any) -> np.ndarray:
+        """Apply ToRGB to a single volume.
+
+        Args:
+            volume (np.ndarray): Volume with shape (D, H, W, C) or (D, H, W).
+            **params (Any): Additional parameters.
+
+        Returns:
+            np.ndarray: Grayscale volume.
+
+        """
+        return self.apply(volume, **params)
+
+    def apply_to_volumes(self, volumes: np.ndarray, **params: Any) -> np.ndarray:
+        """Apply ToRGB to a batch of volumes.
+
+        Args:
+            volumes (np.ndarray): Batch of volumes with shape (N, D, H, W, C) or (N, D, H, W).
+            **params (Any): Additional parameters.
+
+        Returns:
+            np.ndarray: Batch of RGB volumes.
+
+        """
+        return self.apply(volumes, **params)
+
 
 class ToSepia(ImageOnlyTransform):
     """Apply a sepia filter to the input image.

pyproject.toml

@@ -222,6 +222,7 @@ lint.ignore = [
   "FBT002",
   "FBT003",
   "G004",
+  "PLC0415",
   "PLR0911",
   "PLR0913",
   "PLR2004",