Skip to content

Functional Image Operations¤

Low-level functional operations for image processing.

See Also¤

datarax.operators.modality.image.functional ¤

JAX-native image manipulation operations.

This module provides JAX-native implementations of common image manipulation operations, optimized for use with JAX transformations like jit and vmap. All functions are designed to be compatible with JAX's functional programming model.

logger module-attribute ¤

logger = getLogger(__name__)

resize_image_to_shape ¤

resize_image_to_shape(image: Array, output_size: tuple[int, int], method: str = 'bilinear', antialias: bool = True) -> Array

Resize an image to the specified dimensions using JAX-native operations.

Parameters:

Name Type Description Default
image Array

Input image as JAX array with shape [height, width, channels] or [height, width] for grayscale.

 required
output_size tuple[int, int]

Target size as (height, width) tuple.

 required
method str

Interpolation method, one of "nearest", "bilinear", or "bicubic".

 'bilinear'
antialias bool

Whether to use antialiasing when downsampling.

 True

Returns:

Type Description
Array

Resized image with shape [new_height, new_width, channels] or
Array

[new_height, new_width] for grayscale.

center_crop ¤

center_crop(image: Array, output_size: tuple[int, int]) -> Array

Crop the center of an image to the specified dimensions.

Parameters:

Name Type Description Default
image Array

Input image as JAX array with shape [height, width, channels] or [height, width] for grayscale.

 required
output_size tuple[int, int]

Target size as (height, width) tuple.

 required

Returns:

Type Description
Array

Center-cropped image.

random_crop ¤

random_crop(image: Array, output_size: tuple[int, int], key: Array) -> Array

Randomly crop an image to the specified dimensions.

Parameters:

Name Type Description Default
image Array

Input image as JAX array with shape [height, width, channels] or [height, width] for grayscale.

 required
output_size tuple[int, int]

Target size as (height, width) tuple.

 required
key Array

JAX PRNG key for random operations.

 required

Returns:

Type Description
Array

Randomly cropped image.

normalize ¤

normalize(image: Array, mean: float | Array, std: float | Array, clip: bool = False) -> Array

Normalize an image by subtracting mean and dividing by std.

Parameters:

Name Type Description Default
image Array

Input image as JAX array.

 required
mean float | Array

Mean value(s) to subtract. Can be a scalar or array matching the channel dimension.

 required
std float | Array

Standard deviation value(s) to divide by. Can be a scalar or array matching the channel dimension.

 required
clip bool

Whether to clip the output to [0, 1] range.

 False

Returns:

Type Description
Array

Normalized image.

random_flip_left_right ¤

random_flip_left_right(image: Array, key: Array, probability: float = 0.5) -> Array

Randomly flip an image horizontally based on the given probability.

Parameters:

Name Type Description Default
image Array

Input image as JAX array.

 required
key Array

JAX PRNG key for randomness.

 required
probability float

Probability of applying the flip (0.0 to 1.0).

 0.5

Returns:

Type Description
Array

The original or flipped image.

random_flip_up_down ¤

random_flip_up_down(image: Array, key: Array, probability: float = 0.5) -> Array

Randomly flip an image vertically based on the given probability.

Parameters:

Name Type Description Default
image Array

Input image as JAX array.

 required
key Array

JAX PRNG key for randomness.

 required
probability float

Probability of applying the flip (0.0 to 1.0).

 0.5

Returns:

Type Description
Array

The original or flipped image.

adjust_brightness ¤

adjust_brightness(image: Array, factor: float | Array) -> Array

Adjust the brightness of an image.

Parameters:

Name Type Description Default
image Array

Input image as JAX array with values in [0, 1].

 required
factor float | Array

Brightness adjustment factor. Values > 1 increase brightness, < 1 decrease it.

 required

Returns:

Type Description
Array

Brightness-adjusted image, clipped to [0, 1].

adjust_brightness_delta ¤

adjust_brightness_delta(image: Array, delta: float) -> Array

Adjust the brightness of an image using an additive delta.

Parameters:

Name Type Description Default
image Array

Input image as JAX array with values in [0, 1].

 required
delta float

Brightness adjustment delta. Values > 0 increase brightness, < 0 decrease it.

 required

Returns:

Type Description
Array

Brightness-adjusted image, clipped to [0, 1].

adjust_contrast ¤

adjust_contrast(image: Array, factor: float | Array) -> Array

Adjust the contrast of an image.

Parameters:

Name Type Description Default
image Array

Input image as JAX array with values in [0, 1].

 required
factor float | Array

Contrast adjustment factor. Values > 1 increase contrast, < 1 decrease it.

 required

Returns:

Type Description
Array

Contrast-adjusted image, clipped to [0, 1].

rgb_to_hsv ¤

rgb_to_hsv(rgb: Array) -> Array

Convert RGB image to HSV color space.

Implementation follows the algorithm from OpenCV's cvtColor.

Parameters:

Name Type Description Default
rgb Array

RGB image with values in range [0, 1] and shape [..., 3]

 required

Returns:

Type Description
Array

HSV image with values in ranges H: [0, 2π], S: [0, 1], V: [0, 1]

hsv_to_rgb ¤

hsv_to_rgb(hsv: Array) -> Array

Convert HSV image to RGB color space.

Parameters:

Name Type Description Default
hsv Array

HSV image with values in ranges H: [0, 2π], S: [0, 1], V: [0, 1]

 required

Returns:

Type Description
Array

RGB image with values in range [0, 1]

adjust_saturation ¤

adjust_saturation(image: Array, factor: float | Array) -> Array

Adjust the saturation of an image.

Parameters:

Name Type Description Default
image Array

Input RGB image as JAX array with values in [0, 1].

 required
factor float | Array

Saturation adjustment factor. Values > 1 increase saturation, < 1 decrease it.

 required

Returns:

Type Description
Array

Saturation-adjusted image, clipped to [0, 1].

adjust_hue ¤

adjust_hue(image: Array, delta: float | Array) -> Array

Adjust the hue of an image.

Parameters:

Name Type Description Default
image Array

Input RGB image as JAX array with values in [0, 1].

 required
delta float | Array

Hue adjustment in radians. Valid range is [-π, π].

 required

Returns:

Type Description
Array

Hue-adjusted image, clipped to [0, 1].

color_jitter ¤

color_jitter(image: Array, brightness: float = 0.0, contrast: float = 0.0, saturation: float = 0.0, hue: float = 0.0, key: Array | None = None) -> Array

Apply color jittering to an image.

Parameters:

Name Type Description Default
image Array

Input RGB image as JAX array with values in [0, 1].

 required
brightness float

Maximum brightness adjustment factor.

 0.0
contrast float

Maximum contrast adjustment factor.

 0.0
saturation float

Maximum saturation adjustment factor.

 0.0
hue float

Maximum hue adjustment in radians.

 0.0
key Array | None

JAX PRNG key for randomness. If provided, random adjustments are made.

 None

Returns:

Type Description
Array

Color-jittered image, clipped to [0, 1].

convert_rgb_to_grayscale ¤

convert_rgb_to_grayscale(image: Array) -> Array

Convert an RGB image to grayscale.

Parameters:

Name Type Description Default
image Array

Input RGB image as JAX array with shape [height, width, 3].

 required

Returns:

Type Description
Array

Grayscale image with shape [height, width].

rotate ¤

rotate(image: Array, angle_rad: float, fill_value: float = 0.0) -> Array

Rotate image using bilinear interpolation.

This function implements rotation using the inverse transformation approach: 1. For each output pixel, compute itssource location in the input image 2. Use bilinear interpolation to sample the input image at that location 3. Fill empty areas (outside input bounds) with fill_value

Parameters:

Name Type Description Default
image Array

Input image array, shape (H, W, C) or (H, W).

 required
angle_rad float

Rotation angle in radians (counter-clockwise).

 required
fill_value float

Value to fill empty areas after rotation.

 0.0

Returns:

Type Description
Array

Rotated image with same shape as input.