import functools import jax import jax.numpy as jnp from keras.src.backend.jax.core import convert_to_tensor RESIZE_INTERPOLATIONS = ( "bilinear", "nearest", "lanczos3", "lanczos5", "bicubic", ) def rgb_to_grayscale(image, data_format="channels_last"): if data_format == "channels_first": if len(image.shape) == 4: image = jnp.transpose(image, (0, 2, 3, 1)) elif len(image.shape) == 3: image = jnp.transpose(image, (1, 2, 0)) else: raise ValueError( "Invalid input rank: expected rank 3 (single image) " "or rank 4 (batch of images). Received input with shape: " f"image.shape={image.shape}" ) red, green, blue = image[..., 0], image[..., 1], image[..., 2] grayscale_image = 0.2989 * red + 0.5870 * green + 0.1140 * blue grayscale_image = jnp.expand_dims(grayscale_image, axis=-1) if data_format == "channels_first": if len(image.shape) == 4: grayscale_image = jnp.transpose(grayscale_image, (0, 3, 1, 2)) elif len(image.shape) == 3: grayscale_image = jnp.transpose(grayscale_image, (2, 0, 1)) return jnp.array(grayscale_image) def resize( image, size, interpolation="bilinear", antialias=False, crop_to_aspect_ratio=False, pad_to_aspect_ratio=False, fill_mode="constant", fill_value=0.0, data_format="channels_last", ): if interpolation not in RESIZE_INTERPOLATIONS: raise ValueError( "Invalid value for argument `interpolation`. Expected of one " f"{RESIZE_INTERPOLATIONS}. Received: interpolation={interpolation}" ) if fill_mode != "constant": raise ValueError( "Invalid value for argument `fill_mode`. Only `'constant'` " f"is supported. Received: fill_mode={fill_mode}" ) if pad_to_aspect_ratio and crop_to_aspect_ratio: raise ValueError( "Only one of `pad_to_aspect_ratio` & `crop_to_aspect_ratio` " "can be `True`." ) if not len(size) == 2: raise ValueError( "Argument `size` must be a tuple of two elements " f"(height, width). Received: size={size}" ) size = tuple(size) target_height, target_width = size if len(image.shape) == 4: if data_format == "channels_last": size = (image.shape[0],) + size + (image.shape[-1],) else: size = (image.shape[0], image.shape[1]) + size batch_size = image.shape[0] elif len(image.shape) == 3: if data_format == "channels_last": size = size + (image.shape[-1],) else: size = (image.shape[0],) + size else: raise ValueError( "Invalid input rank: expected rank 3 (single image) " "or rank 4 (batch of images). Received input with shape: " f"image.shape={image.shape}" ) if crop_to_aspect_ratio: shape = image.shape if data_format == "channels_last": height, width = shape[-3], shape[-2] else: height, width = shape[-2], shape[-1] crop_height = int(float(width * target_height) / target_width) crop_height = min(height, crop_height) crop_width = int(float(height * target_width) / target_height) crop_width = min(width, crop_width) crop_box_hstart = int(float(height - crop_height) / 2) crop_box_wstart = int(float(width - crop_width) / 2) if data_format == "channels_last": if len(image.shape) == 4: image = image[ :, crop_box_hstart : crop_box_hstart + crop_height, crop_box_wstart : crop_box_wstart + crop_width, :, ] else: image = image[ crop_box_hstart : crop_box_hstart + crop_height, crop_box_wstart : crop_box_wstart + crop_width, :, ] else: if len(image.shape) == 4: image = image[ :, :, crop_box_hstart : crop_box_hstart + crop_height, crop_box_wstart : crop_box_wstart + crop_width, ] else: image = image[ :, crop_box_hstart : crop_box_hstart + crop_height, crop_box_wstart : crop_box_wstart + crop_width, ] elif pad_to_aspect_ratio: shape = image.shape if data_format == "channels_last": height, width, channels = shape[-3], shape[-2], shape[-1] else: height, width, channels = shape[-2], shape[-1], shape[-3] pad_height = int(float(width * target_height) / target_width) pad_height = max(height, pad_height) pad_width = int(float(height * target_width) / target_height) pad_width = max(width, pad_width) img_box_hstart = int(float(pad_height - height) / 2) img_box_wstart = int(float(pad_width - width) / 2) if data_format == "channels_last": if img_box_hstart > 0: if len(image.shape) == 4: padded_img = jnp.concatenate( [ jnp.ones( (batch_size, img_box_hstart, width, channels), dtype=image.dtype, ) * fill_value, image, jnp.ones( (batch_size, img_box_hstart, width, channels), dtype=image.dtype, ) * fill_value, ], axis=1, ) else: padded_img = jnp.concatenate( [ jnp.ones( (img_box_hstart, width, channels), dtype=image.dtype, ) * fill_value, image, jnp.ones( (img_box_hstart, width, channels), dtype=image.dtype, ) * fill_value, ], axis=0, ) elif img_box_wstart > 0: if len(image.shape) == 4: padded_img = jnp.concatenate( [ jnp.ones( (batch_size, height, img_box_wstart, channels), dtype=image.dtype, ) * fill_value, image, jnp.ones( (batch_size, height, img_box_wstart, channels), dtype=image.dtype, ) * fill_value, ], axis=2, ) else: padded_img = jnp.concatenate( [ jnp.ones( (height, img_box_wstart, channels), dtype=image.dtype, ) * fill_value, image, jnp.ones( (height, img_box_wstart, channels), dtype=image.dtype, ) * fill_value, ], axis=1, ) else: padded_img = image else: if img_box_hstart > 0: if len(image.shape) == 4: padded_img = jnp.concatenate( [ jnp.ones( (batch_size, channels, img_box_hstart, width) ) * fill_value, image, jnp.ones( (batch_size, channels, img_box_hstart, width) ) * fill_value, ], axis=2, ) else: padded_img = jnp.concatenate( [ jnp.ones((channels, img_box_hstart, width)) * fill_value, image, jnp.ones((channels, img_box_hstart, width)) * fill_value, ], axis=1, ) elif img_box_wstart > 0: if len(image.shape) == 4: padded_img = jnp.concatenate( [ jnp.ones( (batch_size, channels, height, img_box_wstart) ) * fill_value, image, jnp.ones( (batch_size, channels, height, img_box_wstart) ) * fill_value, ], axis=3, ) else: padded_img = jnp.concatenate( [ jnp.ones((channels, height, img_box_wstart)) * fill_value, image, jnp.ones((channels, height, img_box_wstart)) * fill_value, ], axis=2, ) else: padded_img = image image = padded_img return jax.image.resize( image, size, method=interpolation, antialias=antialias ) AFFINE_TRANSFORM_INTERPOLATIONS = { # map to order "nearest": 0, "bilinear": 1, } AFFINE_TRANSFORM_FILL_MODES = { "constant", "nearest", "wrap", "mirror", "reflect", } def affine_transform( image, transform, interpolation="bilinear", fill_mode="constant", fill_value=0, data_format="channels_last", ): if interpolation not in AFFINE_TRANSFORM_INTERPOLATIONS.keys(): raise ValueError( "Invalid value for argument `interpolation`. Expected of one " f"{set(AFFINE_TRANSFORM_INTERPOLATIONS.keys())}. Received: " f"interpolation={interpolation}" ) if fill_mode not in AFFINE_TRANSFORM_FILL_MODES: raise ValueError( "Invalid value for argument `fill_mode`. Expected of one " f"{AFFINE_TRANSFORM_FILL_MODES}. Received: fill_mode={fill_mode}" ) transform = convert_to_tensor(transform) if len(image.shape) not in (3, 4): raise ValueError( "Invalid image rank: expected rank 3 (single image) " "or rank 4 (batch of images). Received input with shape: " f"image.shape={image.shape}" ) if len(transform.shape) not in (1, 2): raise ValueError( "Invalid transform rank: expected rank 1 (single transform) " "or rank 2 (batch of transforms). Received input with shape: " f"transform.shape={transform.shape}" ) # unbatched case need_squeeze = False if len(image.shape) == 3: image = jnp.expand_dims(image, axis=0) need_squeeze = True if len(transform.shape) == 1: transform = jnp.expand_dims(transform, axis=0) if data_format == "channels_first": image = jnp.transpose(image, (0, 2, 3, 1)) batch_size = image.shape[0] # get indices meshgrid = jnp.meshgrid( *[jnp.arange(size) for size in image.shape[1:]], indexing="ij" ) indices = jnp.concatenate( [jnp.expand_dims(x, axis=-1) for x in meshgrid], axis=-1 ) indices = jnp.tile(indices, (batch_size, 1, 1, 1, 1)) # swap the values a0 = transform[:, 0] a2 = transform[:, 2] b1 = transform[:, 4] b2 = transform[:, 5] transform = transform.at[:, 0].set(b1) transform = transform.at[:, 2].set(b2) transform = transform.at[:, 4].set(a0) transform = transform.at[:, 5].set(a2) # deal with transform transform = jnp.pad( transform, pad_width=[[0, 0], [0, 1]], constant_values=1 ) transform = jnp.reshape(transform, (batch_size, 3, 3)) offset = transform[:, 0:2, 2] offset = jnp.pad(offset, pad_width=[[0, 0], [0, 1]]) transform = transform.at[:, 0:2, 2].set(0) # transform the indices coordinates = jnp.einsum("Bhwij, Bjk -> Bhwik", indices, transform) coordinates = jnp.moveaxis(coordinates, source=-1, destination=1) coordinates += jnp.reshape(a=offset, newshape=(*offset.shape, 1, 1, 1)) # apply affine transformation _map_coordinates = functools.partial( jax.scipy.ndimage.map_coordinates, order=AFFINE_TRANSFORM_INTERPOLATIONS[interpolation], mode=fill_mode, cval=fill_value, ) affined = jax.vmap(_map_coordinates)(image, coordinates) if data_format == "channels_first": affined = jnp.transpose(affined, (0, 3, 1, 2)) if need_squeeze: affined = jnp.squeeze(affined, axis=0) return affined MAP_COORDINATES_FILL_MODES = { "constant", "nearest", "wrap", "mirror", "reflect", } def map_coordinates( input, coordinates, order, fill_mode="constant", fill_value=0.0 ): if fill_mode not in MAP_COORDINATES_FILL_MODES: raise ValueError( "Invalid value for argument `fill_mode`. Expected one of " f"{set(MAP_COORDINATES_FILL_MODES)}. Received: " f"fill_mode={fill_mode}" ) if order not in range(2): raise ValueError( "Invalid value for argument `order`. Expected one of " f"{[0, 1]}. Received: order={order}" ) return jax.scipy.ndimage.map_coordinates( input, coordinates, order, fill_mode, fill_value )