Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Update posterize #2180

Merged
merged 2 commits into from
Dec 7, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
49 changes: 33 additions & 16 deletions albumentations/augmentations/functional.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -167,23 +166,44 @@

@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
Expand All @@ -192,14 +212,11 @@

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

Check warning on line 217 in albumentations/augmentations/functional.py

View check run for this annotation

Codecov / codecov/patch

albumentations/augmentations/functional.py#L215-L217

Added lines #L215 - L217 were not covered by tests

result_img[..., i] = sz_lut(img[..., i], lut, inplace=True)
result_img[..., i] = sz_lut(img[..., i], lut, inplace=True)

Check warning on line 219 in albumentations/augmentations/functional.py

View check run for this annotation

Codecov / codecov/patch

albumentations/augmentations/functional.py#L219

Added line #L219 was not covered by tests

return result_img

Expand Down
19 changes: 11 additions & 8 deletions albumentations/augmentations/transforms.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -2080,8 +2080,8 @@
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
Expand All @@ -2099,8 +2099,6 @@

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
Expand Down Expand Up @@ -2149,8 +2147,8 @@
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]")

Check warning on line 2151 in albumentations/augmentations/transforms.py

View check run for this annotation

Codecov / codecov/patch

albumentations/augmentations/transforms.py#L2151

Added line #L2151 was not covered by tests
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]
Expand All @@ -2165,7 +2163,12 @@
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]:
Expand Down
1 change: 1 addition & 0 deletions albumentations/core/types.py
Original file line number Diff line number Diff line change
Expand Up @@ -57,6 +57,7 @@ class Targets(Enum):
TWO = 2
THREE = 3
FOUR = 4
SEVEN = 7
EIGHT = 8
THREE_SIXTY = 360

Expand Down
Loading