nrtk-albumentations 2.1.0__py3-none-any.whl

albumentations/augmentations/spectrogram/transform.py

@@ -0,0 +1,220 @@
+"""Transforms for spectrogram augmentation.
+
+This module provides transforms specifically designed for augmenting spectrograms
+in audio processing tasks. Includes time reversal, time masking, and frequency
+masking transforms commonly used in audio machine learning applications.
+"""
+
+from __future__ import annotations
+
+from warnings import warn
+
+from pydantic import Field
+
+from albumentations.augmentations.dropout.xy_masking import XYMasking
+from albumentations.augmentations.geometric.flip import HorizontalFlip
+from albumentations.core.transforms_interface import BaseTransformInitSchema
+from albumentations.core.type_definitions import ALL_TARGETS
+
+__all__ = [
+    "FrequencyMasking",
+    "TimeMasking",
+    "TimeReverse",
+]
+
+
+class TimeReverse(HorizontalFlip):
+    """Reverse the time axis of a spectrogram image, also known as time inversion.
+
+    Time inversion of a spectrogram is analogous to the random flip of an image,
+    an augmentation technique widely used in the visual domain. This can be relevant
+    in the context of audio classification tasks when working with spectrograms.
+    The technique was successfully applied in the AudioCLIP paper, which extended
+    CLIP to handle image, text, and audio inputs.
+
+    This transform is implemented as a subclass of HorizontalFlip since reversing
+    time in a spectrogram is equivalent to flipping the image horizontally.
+
+    Args:
+        p (float): probability of applying the transform. Default: 0.5.
+
+    Targets:
+        image, mask, bboxes, keypoints, volume, mask3d
+
+    Image types:
+        uint8, float32
+
+    Number of channels:
+        Any
+
+    Note:
+        This transform is functionally identical to HorizontalFlip but provides
+        a more semantically meaningful name when working with spectrograms and
+        other time-series visualizations.
+
+    References:
+        - AudioCLIP paper: https://arxiv.org/abs/2106.13043
+        - Audiomentations: https://iver56.github.io/audiomentations/waveform_transforms/reverse/
+
+    """
+
+    _targets = ALL_TARGETS
+
+    class InitSchema(BaseTransformInitSchema):
+        pass
+
+    def __init__(
+        self,
+        p: float = 0.5,
+    ):
+        warn(
+            "TimeReverse is an alias for HorizontalFlip transform. "
+            "Consider using HorizontalFlip directly from albumentations.HorizontalFlip. ",
+            UserWarning,
+            stacklevel=2,
+        )
+        super().__init__(p=p)
+
+
+class TimeMasking(XYMasking):
+    """Apply masking to a spectrogram in the time domain.
+
+    This transform masks random segments along the time axis of a spectrogram,
+    implementing the time masking technique proposed in the SpecAugment paper.
+    Time masking helps in training models to be robust against temporal variations
+    and missing information in audio signals.
+
+    This is a specialized version of XYMasking configured for time masking only.
+    For more advanced use cases (e.g., multiple masks, frequency masking, or custom
+    fill values), consider using XYMasking directly.
+
+    Args:
+        time_mask_param (int): Maximum possible length of the mask in the time domain.
+            Must be a positive integer. Length of the mask is uniformly sampled from (0, time_mask_param).
+        p (float): probability of applying the transform. Default: 0.5.
+
+    Targets:
+        image, mask, bboxes, keypoints, volume, mask3d
+
+    Image types:
+        uint8, float32
+
+    Number of channels:
+        Any
+
+    Note:
+        This transform is implemented as a subset of XYMasking with fixed parameters:
+        - Single horizontal mask (num_masks_x=1)
+        - No vertical masks (num_masks_y=0)
+        - Zero fill value
+        - Random mask length up to time_mask_param
+
+        For more flexibility, including:
+        - Multiple masks
+        - Custom fill values
+        - Frequency masking
+        - Combined time-frequency masking
+        Consider using albumentations.XYMasking directly.
+
+    References:
+        - SpecAugment paper: https://arxiv.org/abs/1904.08779
+        - Original implementation: https://pytorch.org/audio/stable/transforms.html#timemask
+
+    """
+
+    class InitSchema(BaseTransformInitSchema):
+        time_mask_param: int = Field(gt=0)
+
+    def __init__(
+        self,
+        time_mask_param: int = 40,
+        p: float = 0.5,
+    ):
+        warn(
+            "TimeMasking is a specialized version of XYMasking. "
+            "For more flexibility (multiple masks, custom fill values, frequency masking), "
+            "consider using XYMasking directly from albumentations.XYMasking.",
+            UserWarning,
+            stacklevel=2,
+        )
+        super().__init__(
+            num_masks_x=1,
+            num_masks_y=0,
+            mask_x_length=(0, time_mask_param),
+            fill=0,
+            fill_mask=0,
+            p=p,
+        )
+        self.time_mask_param = time_mask_param
+
+
+class FrequencyMasking(XYMasking):
+    """Apply masking to a spectrogram in the frequency domain.
+
+    This transform masks random segments along the frequency axis of a spectrogram,
+    implementing the frequency masking technique proposed in the SpecAugment paper.
+    Frequency masking helps in training models to be robust against frequency variations
+    and missing spectral information in audio signals.
+
+    This is a specialized version of XYMasking configured for frequency masking only.
+    For more advanced use cases (e.g., multiple masks, time masking, or custom
+    fill values), consider using XYMasking directly.
+
+    Args:
+        freq_mask_param (int): Maximum possible length of the mask in the frequency domain.
+            Must be a positive integer. Length of the mask is uniformly sampled from (0, freq_mask_param).
+        p (float): probability of applying the transform. Default: 0.5.
+
+    Targets:
+        image, mask, bboxes, keypoints, volume, mask3d
+
+    Image types:
+        uint8, float32
+
+    Number of channels:
+        Any
+
+    Note:
+        This transform is implemented as a subset of XYMasking with fixed parameters:
+        - Single vertical mask (num_masks_y=1)
+        - No horizontal masks (num_masks_x=0)
+        - Zero fill value
+        - Random mask length up to freq_mask_param
+
+        For more flexibility, including:
+        - Multiple masks
+        - Custom fill values
+        - Time masking
+        - Combined time-frequency masking
+        Consider using albumentations.XYMasking directly.
+
+    References:
+        - SpecAugment paper: https://arxiv.org/abs/1904.08779
+        - Original implementation: https://pytorch.org/audio/stable/transforms.html#freqmask
+
+    """
+
+    class InitSchema(BaseTransformInitSchema):
+        freq_mask_param: int = Field(gt=0)
+
+    def __init__(
+        self,
+        freq_mask_param: int = 30,
+        p: float = 0.5,
+    ):
+        warn(
+            "FrequencyMasking is a specialized version of XYMasking. "
+            "For more flexibility (multiple masks, custom fill values, time masking), "
+            "consider using XYMasking directly from albumentations.XYMasking.",
+            UserWarning,
+            stacklevel=2,
+        )
+        super().__init__(
+            p=p,
+            fill=0,
+            fill_mask=0,
+            mask_y_length=(0, freq_mask_param),
+            num_masks_x=0,
+            num_masks_y=1,
+        )
+        self.freq_mask_param = freq_mask_param

albumentations/augmentations/text/functional.py

@@ -0,0 +1,272 @@
+"""Functional implementations for text manipulation and rendering.
+
+This module provides utility functions for manipulating text in strings and
+rendering text onto images. Includes functions for word manipulation, text drawing,
+and handling text regions in images.
+"""
+
+from __future__ import annotations
+
+import random
+from typing import TYPE_CHECKING, Any
+
+import cv2
+import numpy as np
+from albucore import (
+    MONO_CHANNEL_DIMENSIONS,
+    NUM_MULTI_CHANNEL_DIMENSIONS,
+    NUM_RGB_CHANNELS,
+    preserve_channel_dim,
+    uint8_io,
+)
+
+from albumentations.core.type_definitions import PAIR
+
+# Importing wordnet and other dependencies only for type checking
+if TYPE_CHECKING:
+    from PIL import Image
+
+
+def delete_random_words(words: list[str], num_words: int, py_random: random.Random) -> str:
+    """Delete a specified number of random words from a list.
+
+    This function randomly removes words from the input list and joins the remaining
+    words with spaces to form a new string.
+
+    Args:
+        words (list[str]): List of words to process.
+        num_words (int): Number of words to delete.
+        py_random (random.Random): Random number generator for reproducibility.
+
+    Returns:
+        str: New string with specified words removed. Returns empty string if
+             num_words is greater than or equal to the length of words.
+
+    """
+    if num_words >= len(words):
+        return ""
+
+    indices_to_delete = py_random.sample(range(len(words)), num_words)
+    new_words = [word for idx, word in enumerate(words) if idx not in indices_to_delete]
+    return " ".join(new_words)
+
+
+def swap_random_words(words: list[str], num_words: int, py_random: random.Random) -> str:
+    """Swap random pairs of words in a list of words.
+
+    This function randomly selects pairs of words and swaps their positions
+    a specified number of times.
+
+    Args:
+        words (list[str]): List of words to process.
+        num_words (int): Number of swaps to perform.
+        py_random (random.Random): Random number generator for reproducibility.
+
+    Returns:
+        str: New string with words swapped. If num_words is 0 or the list has fewer
+             than 2 words, returns the original string.
+
+    """
+    if num_words == 0 or len(words) < PAIR:
+        return " ".join(words)
+
+    words = words.copy()
+
+    for _ in range(num_words):
+        idx1, idx2 = py_random.sample(range(len(words)), 2)
+        words[idx1], words[idx2] = words[idx2], words[idx1]
+    return " ".join(words)
+
+
+def insert_random_stopwords(
+    words: list[str],
+    num_insertions: int,
+    stopwords: tuple[str, ...] | None,
+    py_random: random.Random,
+) -> str:
+    """Insert random stopwords into a list of words.
+
+    This function randomly inserts stopwords at random positions in the
+    list of words a specified number of times.
+
+    Args:
+        words (list[str]): List of words to process.
+        num_insertions (int): Number of stopwords to insert.
+        stopwords (tuple[str, ...] | None): Tuple of stopwords to choose from.
+            If None, default stopwords will be used.
+        py_random (random.Random): Random number generator for reproducibility.
+
+    Returns:
+        str: New string with stopwords inserted.
+
+    """
+    if stopwords is None:
+        stopwords = ("and", "the", "is", "in", "at", "of")  # Default stopwords if none provided
+
+    for _ in range(num_insertions):
+        idx = py_random.randint(0, len(words))
+        words.insert(idx, py_random.choice(stopwords))
+    return " ".join(words)
+
+
+def convert_image_to_pil(image: np.ndarray) -> Image:
+    """Convert a NumPy array image to a PIL image."""
+    try:
+        from PIL import Image
+    except ImportError:
+        raise ImportError("Pillow is not installed") from ImportError
+
+    if image.ndim == MONO_CHANNEL_DIMENSIONS:  # (height, width)
+        return Image.fromarray(image)
+    if image.ndim == NUM_MULTI_CHANNEL_DIMENSIONS and image.shape[2] == 1:  # (height, width, 1)
+        return Image.fromarray(image[:, :, 0], mode="L")
+    if image.ndim == NUM_MULTI_CHANNEL_DIMENSIONS and image.shape[2] == NUM_RGB_CHANNELS:  # (height, width, 3)
+        return Image.fromarray(image)
+
+    raise TypeError(f"Unsupported image shape: {image.shape}")
+
+
+def draw_text_on_pil_image(pil_image: Image, metadata_list: list[dict[str, Any]]) -> Image:
+    """Draw text on a PIL image."""
+    try:
+        from PIL import ImageDraw
+    except ImportError:
+        raise ImportError("Pillow is not installed") from ImportError
+
+    draw = ImageDraw.Draw(pil_image)
+    for metadata in metadata_list:
+        bbox_coords = metadata["bbox_coords"]
+        text = metadata["text"]
+        font = metadata["font"]
+        font_color = metadata["font_color"]
+
+        # Adapt font_color based on image mode
+        if pil_image.mode == "L":  # Grayscale
+            # For grayscale images, use only the first value or average the RGB values
+            if isinstance(font_color, tuple):
+                if len(font_color) >= 3:
+                    # Average RGB values for grayscale
+                    font_color = int(sum(font_color[:3]) / 3)
+                elif len(font_color) == 1:
+                    font_color = int(font_color[0])
+        # For RGB and other modes, ensure font_color is a tuple of integers
+        elif isinstance(font_color, tuple):
+            font_color = tuple(int(c) for c in font_color)
+
+        position = bbox_coords[:2]
+        draw.text(position, text, font=font, fill=font_color)
+    return pil_image
+
+
+def draw_text_on_multi_channel_image(image: np.ndarray, metadata_list: list[dict[str, Any]]) -> np.ndarray:
+    """Draw text on a multi-channel image with more than three channels."""
+    try:
+        from PIL import Image, ImageDraw
+    except ImportError:
+        raise ImportError("Pillow is not installed") from ImportError
+
+    channels = [Image.fromarray(image[:, :, i]) for i in range(image.shape[2])]
+    pil_images = [ImageDraw.Draw(channel) for channel in channels]
+
+    for metadata in metadata_list:
+        bbox_coords = metadata["bbox_coords"]
+        text = metadata["text"]
+        font = metadata["font"]
+        font_color = metadata["font_color"]
+
+        # Handle font_color as tuple[float, ...]
+        # Ensure we have enough color values for all channels
+        if len(font_color) < image.shape[2]:
+            # If fewer values than channels, pad with zeros
+            font_color = tuple(list(font_color) + [0] * (image.shape[2] - len(font_color)))
+        elif len(font_color) > image.shape[2]:
+            # If more values than channels, truncate
+            font_color = font_color[: image.shape[2]]
+
+        # Convert to integers for PIL
+        font_color = [int(c) for c in font_color]
+
+        position = bbox_coords[:2]
+
+        # For each channel, use the corresponding color value
+        for channel_id, pil_image in enumerate(pil_images):
+            # For single-channel PIL images, color must be an integer
+            pil_image.text(position, text, font=font, fill=font_color[channel_id])
+
+    return np.stack([np.array(channel) for channel in channels], axis=2)
+
+
+@uint8_io
+@preserve_channel_dim
+def render_text(image: np.ndarray, metadata_list: list[dict[str, Any]], clear_bg: bool) -> np.ndarray:
+    """Render text onto an image based on provided metadata.
+
+    This function draws text on an image using metadata that specifies text content,
+    position, font, and color. It can optionally clear the background before rendering.
+    The function handles different image types (grayscale, RGB, multi-channel).
+
+    Args:
+        image (np.ndarray): Image to draw text on.
+        metadata_list (list[dict[str, Any]]): List of metadata dictionaries containing:
+            - bbox_coords: Bounding box coordinates (x_min, y_min, x_max, y_max)
+            - text: Text string to render
+            - font: PIL ImageFont object
+            - font_color: Color for the text
+        clear_bg (bool): Whether to clear (inpaint) the background under the text.
+
+    Returns:
+        np.ndarray: Image with text rendered on it.
+
+    """
+    # First clean background under boxes using seamless clone if clear_bg is True
+    if clear_bg:
+        image = inpaint_text_background(image, metadata_list)
+
+    if len(image.shape) == MONO_CHANNEL_DIMENSIONS or (
+        len(image.shape) == NUM_MULTI_CHANNEL_DIMENSIONS and image.shape[2] in {1, NUM_RGB_CHANNELS}
+    ):
+        pil_image = convert_image_to_pil(image)
+        pil_image = draw_text_on_pil_image(pil_image, metadata_list)
+        return np.array(pil_image)
+
+    return draw_text_on_multi_channel_image(image, metadata_list)
+
+
+def inpaint_text_background(
+    image: np.ndarray,
+    metadata_list: list[dict[str, Any]],
+    method: int = cv2.INPAINT_TELEA,
+) -> np.ndarray:
+    """Inpaint (clear) regions in an image where text will be rendered.
+
+    This function creates a clean background for text by inpainting rectangular
+    regions specified in the metadata. It removes any existing content in those
+    regions to provide a clean slate for rendering text.
+
+    Args:
+        image (np.ndarray): Image to inpaint.
+        metadata_list (list[dict[str, Any]]): List of metadata dictionaries containing:
+            - bbox_coords: Bounding box coordinates (x_min, y_min, x_max, y_max)
+        method (int, optional): Inpainting method to use. Defaults to cv2.INPAINT_TELEA.
+            Options include:
+            - cv2.INPAINT_TELEA: Fast Marching Method
+            - cv2.INPAINT_NS: Navier-Stokes method
+
+    Returns:
+        np.ndarray: Image with specified regions inpainted.
+
+    """
+    result_image = image.copy()
+    mask = np.zeros((image.shape[0], image.shape[1]), dtype=np.uint8)
+
+    for metadata in metadata_list:
+        x_min, y_min, x_max, y_max = metadata["bbox_coords"]
+
+        # Black out the region
+        result_image[y_min:y_max, x_min:x_max] = 0
+
+        # Update the mask to indicate the region to inpaint
+        mask[y_min:y_max, x_min:x_max] = 255
+
+    # Inpaint the blacked-out regions
+    return cv2.inpaint(result_image, mask, inpaintRadius=3, flags=method)