diff --git a/albumentations/augmentations/functional.py b/albumentations/augmentations/functional.py index 75c37031c..52adf80df 100644 --- a/albumentations/augmentations/functional.py +++ b/albumentations/augmentations/functional.py @@ -39,7 +39,6 @@ ) from albumentations.core.bbox_utils import bboxes_from_masks, masks_from_bboxes from albumentations.core.types import ( - EIGHT, MONO_CHANNEL_DIMENSIONS, NUM_MULTI_CHANNEL_DIMENSIONS, NUM_RGB_CHANNELS, @@ -167,23 +166,44 @@ def solarize(img: np.ndarray, threshold: float) -> np.ndarray: @uint8_io @clipped -def posterize(img: np.ndarray, bits: Literal[1, 2, 3, 4, 5, 6, 7, 8]) -> np.ndarray: - """Reduce the number of bits for each color channel. +def posterize(img: np.ndarray, bits: Literal[1, 2, 3, 4, 5, 6, 7] | list[Literal[1, 2, 3, 4, 5, 6, 7]]) -> np.ndarray: + """Reduce the number of bits for each color channel by keeping only the highest N bits. + + This transform performs bit-depth reduction by masking out lower bits, effectively + reducing the number of possible values per channel. This creates a posterization + effect where similar colors are merged together. Args: - img: image to posterize. - bits: number of high bits. Must be in range [1, 8] + img: Input image. Can be single or multi-channel. + bits: Number of high bits to keep. Must be in range [1, 7]. + Can be either: + - A single value to apply the same bit reduction to all channels + - A list of values to apply different bit reduction per channel. + Length of list must match number of channels in image. Returns: - Image with reduced color channels. + np.ndarray: Image with reduced bit depth. Has same shape and dtype as input. + + Note: + - The transform keeps the N highest bits and sets all other bits to 0 + - For example, if bits=3: + - Original value: 11010110 (214) + - Keep 3 bits: 11000000 (192) + - The number of unique colors per channel will be 2^bits + - Higher bits values = more colors = more subtle effect + - Lower bits values = fewer colors = more dramatic posterization + Examples: + >>> import numpy as np + >>> image = np.random.randint(0, 256, (100, 100, 3), dtype=np.uint8) + >>> # Same posterization for all channels + >>> result = posterize(image, bits=3) + >>> # Different posterization per channel + >>> result = posterize(image, bits=[3, 4, 5]) # RGB channels """ bits_array = np.uint8(bits) if not bits_array.shape or len(bits_array) == 1: - if bits_array == EIGHT: - return img - lut = np.arange(0, 256, dtype=np.uint8) mask = ~np.uint8(2 ** (8 - bits_array) - 1) lut &= mask @@ -192,14 +212,11 @@ def posterize(img: np.ndarray, bits: Literal[1, 2, 3, 4, 5, 6, 7, 8]) -> np.ndar result_img = np.empty_like(img) for i, channel_bits in enumerate(bits_array): - if channel_bits == EIGHT: - result_img[..., i] = img[..., i].copy() - else: - lut = np.arange(0, 256, dtype=np.uint8) - mask = ~np.uint8(2 ** (8 - channel_bits) - 1) - lut &= mask + lut = np.arange(0, 256, dtype=np.uint8) + mask = ~np.uint8(2 ** (8 - channel_bits) - 1) + lut &= mask - result_img[..., i] = sz_lut(img[..., i], lut, inplace=True) + result_img[..., i] = sz_lut(img[..., i], lut, inplace=True) return result_img diff --git a/albumentations/augmentations/transforms.py b/albumentations/augmentations/transforms.py index 5565c266a..f4fc56f93 100644 --- a/albumentations/augmentations/transforms.py +++ b/albumentations/augmentations/transforms.py @@ -70,11 +70,11 @@ NoOp, ) from albumentations.core.types import ( - EIGHT, MAX_RAIN_ANGLE, MONO_CHANNEL_DIMENSIONS, NUM_RGB_CHANNELS, PAIR, + SEVEN, ChromaticAberrationMode, ColorType, ImageMode, @@ -2080,8 +2080,8 @@ class Posterize(ImageOnlyTransform): Args: num_bits (int | tuple[int, int] | list[int] | list[tuple[int, int]]): Defines the number of bits to keep for each color channel. Can be specified in several ways: - - Single int: Same number of bits for all channels. Range: [1, 8]. - - tuple of two ints: (min_bits, max_bits) to randomly choose from. Range for each: [1, 8]. + - Single int: Same number of bits for all channels. Range: [1, 7]. + - tuple of two ints: (min_bits, max_bits) to randomly choose from. Range for each: [1, 7]. - list of three ints: Specific number of bits for each channel [r_bits, g_bits, b_bits]. - list of three tuples: Ranges for each channel [(r_min, r_max), (g_min, g_max), (b_min, b_max)]. Default: 4 @@ -2099,8 +2099,6 @@ class Posterize(ImageOnlyTransform): Note: - The effect becomes more pronounced as the number of bits is reduced. - - Using 0 bits for a channel will reduce it to a single color (usually black). - - Using 8 bits leaves the channel unchanged. - This transform can create interesting artistic effects or be used for image compression simulation. - Posterization is particularly useful for: * Creating stylized or retro-looking images @@ -2149,8 +2147,8 @@ def validate_num_bits( num_bits: Any, ) -> tuple[int, int] | list[tuple[int, int]]: if isinstance(num_bits, int): - if num_bits < 1 or num_bits > EIGHT: - raise ValueError("num_bits must be in the range [1, 8]") + if num_bits < 1 or num_bits > SEVEN: + raise ValueError("num_bits must be in the range [1, 7]") return (num_bits, num_bits) if isinstance(num_bits, Sequence) and len(num_bits) > PAIR: return [to_tuple(i, i) for i in num_bits] @@ -2165,7 +2163,12 @@ def __init__( super().__init__(p=p, always_apply=always_apply) self.num_bits = cast(Union[tuple[int, int], list[tuple[int, int]]], num_bits) - def apply(self, img: np.ndarray, num_bits: int, **params: Any) -> np.ndarray: + def apply( + self, + img: np.ndarray, + num_bits: Literal[1, 2, 3, 4, 5, 6, 7] | list[Literal[1, 2, 3, 4, 5, 6, 7]], + **params: Any, + ) -> np.ndarray: return fmain.posterize(img, num_bits) def get_params(self) -> dict[str, Any]: diff --git a/albumentations/core/types.py b/albumentations/core/types.py index 8ad182f90..e14878c89 100644 --- a/albumentations/core/types.py +++ b/albumentations/core/types.py @@ -57,6 +57,7 @@ class Targets(Enum): TWO = 2 THREE = 3 FOUR = 4 +SEVEN = 7 EIGHT = 8 THREE_SIXTY = 360