From 9bb680a96a3c07cd1de8735ef30424d513683ce6 Mon Sep 17 00:00:00 2001 From: Hana Joo Date: Tue, 7 Jul 2026 17:38:26 -0700 Subject: [PATCH] Automated Code Change PiperOrigin-RevId: 944182032 --- dm_pix/__init__.py | 2 +- dm_pix/_src/augment.py | 36 ++++++++++++++-------------- dm_pix/_src/color_conversion.py | 14 +++++------ dm_pix/_src/color_conversion_test.py | 6 ++--- dm_pix/_src/interpolation.py | 19 ++++++++------- dm_pix/_src/metrics.py | 10 ++++---- examples/image_augmentation.py | 12 +++++----- 7 files changed, 50 insertions(+), 49 deletions(-) diff --git a/dm_pix/__init__.py b/dm_pix/__init__.py index 064006b..fe21ebc 100644 --- a/dm_pix/__init__.py +++ b/dm_pix/__init__.py @@ -131,6 +131,6 @@ # || || # try: - del _src # pylint: disable=undefined-variable + del _src # pylint: disable=undefined-variable # pyrefly: ignore[unbound-name] except NameError: pass diff --git a/dm_pix/_src/augment.py b/dm_pix/_src/augment.py index 4a50e8a..adad96f 100644 --- a/dm_pix/_src/augment.py +++ b/dm_pix/_src/augment.py @@ -131,7 +131,7 @@ def adjust_hue( rgb = color_conversion.split_channels(image, channel_axis) hue, saturation, value = color_conversion.rgb_planes_to_hsv_planes(*rgb) - rgb_adjusted = color_conversion.hsv_planes_to_rgb_planes((hue + delta) % 1.0, + rgb_adjusted = color_conversion.hsv_planes_to_rgb_planes((hue + delta) % 1.0, # pyrefly: ignore[bad-argument-type, unsupported-operation] saturation, value) return jnp.stack(rgb_adjusted, axis=channel_axis) @@ -225,11 +225,11 @@ def elastic_deformation( shift_map_i = gaussian_blur( image=noise_i, sigma=sigma, - kernel_size=kernel_size) * alpha + kernel_size=kernel_size) * alpha # pyrefly: ignore[bad-argument-type] shift_map_j = gaussian_blur( image=noise_j, sigma=sigma, - kernel_size=kernel_size) * alpha + kernel_size=kernel_size) * alpha # pyrefly: ignore[bad-argument-type] meshgrid = list( jnp.meshgrid( @@ -244,7 +244,7 @@ def elastic_deformation( ) transformed_image = jnp.concatenate([ interpolate_function( - image[..., channel, jnp.newaxis], jnp.asarray(meshgrid)) + image[..., channel, jnp.newaxis], jnp.asarray(meshgrid)) # pyrefly: ignore[bad-index] for channel in range(image.shape[-1]) ], axis=-1) @@ -287,9 +287,9 @@ def center_crop( ) center_h, center_w = current_height // 2, current_width // 2 - left = max(center_w - (width // 2), 0) + left = max(center_w - (width // 2), 0) # pyrefly: ignore[unsupported-operation] right = min(left + width, current_width) - top = max(center_h - (height // 2), 0) + top = max(center_h - (height // 2), 0) # pyrefly: ignore[unsupported-operation] bottom = min(top + height, current_height) if _channels_last(image, channel_axis): @@ -304,7 +304,7 @@ def center_crop( limit_indices = (batch, *limit_indices) return jax.lax.slice( - image, start_indices=start_indices, limit_indices=limit_indices + image, start_indices=start_indices, limit_indices=limit_indices # pyrefly: ignore[bad-argument-type] ) @@ -407,8 +407,8 @@ def resize_with_crop_or_pad( ) return pad_to_size( image, - target_height=target_height, - target_width=target_width, + target_height=target_height, # pyrefly: ignore[bad-argument-type] + target_width=target_width, # pyrefly: ignore[bad-argument-type] channel_axis=channel_axis, mode=pad_mode, pad_kwargs=pad_kwargs, @@ -507,7 +507,7 @@ def gaussian_blur( expand_batch_dim = image.ndim == 3 if expand_batch_dim: - image = image[jnp.newaxis, ...] + image = image[jnp.newaxis, ...] # pyrefly: ignore[bad-index] blurred = _depthwise_conv2d( image, kernel=blur_h, @@ -710,8 +710,8 @@ def affine_transform( [jnp.expand_dims(x, axis=-1) for x in meshgrid], axis=-1) if matrix.shape == (4, 4) or matrix.shape == (3, 4): - offset = matrix[:image.ndim, image.ndim] - matrix = matrix[:image.ndim, :image.ndim] + offset = matrix[:image.ndim, image.ndim] # pyrefly: ignore[bad-index] + matrix = matrix[:image.ndim, :image.ndim] # pyrefly: ignore[bad-index] coordinates = indices @ matrix.T coordinates = jnp.moveaxis(coordinates, source=-1, destination=0) @@ -823,7 +823,7 @@ def random_brightness( """`adjust_brightness(...)` with random delta in `[-max_delta, max_delta)`.""" # DO NOT REMOVE - Logging usage. - delta = jax.random.uniform(key, (), minval=-max_delta, maxval=max_delta) + delta = jax.random.uniform(key, (), minval=-max_delta, maxval=max_delta) # pyrefly: ignore[unsupported-operation] return adjust_brightness(image, delta) @@ -854,7 +854,7 @@ def random_hue( """`adjust_hue(...)` with random delta in `[-max_delta, max_delta)`.""" # DO NOT REMOVE - Logging usage. - delta = jax.random.uniform(key, (), minval=-max_delta, maxval=max_delta) + delta = jax.random.uniform(key, (), minval=-max_delta, maxval=max_delta) # pyrefly: ignore[unsupported-operation] return adjust_hue(image, delta, channel_axis=channel_axis) @@ -917,7 +917,7 @@ def random_crop( assert len(image_shape) == len(crop_sizes), ( f"Number of image dims {len(image_shape)} and number of crop_sizes " f"{len(crop_sizes)} do not match.") - assert image_shape >= crop_sizes, ( + assert image_shape >= crop_sizes, ( # pyrefly: ignore[unsupported-operation] f"Crop sizes {crop_sizes} should be a subset of image size {image_shape} " "in each dimension .") random_keys = jax.random.split(key, len(crop_sizes)) @@ -926,7 +926,7 @@ def random_crop( jax.random.randint(k, (), 0, img_size - crop_size + 1) for k, img_size, crop_size in zip(random_keys, image_shape, crop_sizes) ] - out = jax.lax.dynamic_slice(image, slice_starts, crop_sizes) + out = jax.lax.dynamic_slice(image, slice_starts, crop_sizes) # pyrefly: ignore[bad-argument-type] return out @@ -964,8 +964,8 @@ def _depthwise_conv2d( dimension_numbers as the input. """ return jax.lax.conv_general_dilated( - inputs, - kernel, + inputs, # pyrefly: ignore[bad-argument-type] + kernel, # pyrefly: ignore[bad-argument-type] strides, padding, feature_group_count=inputs.shape[channel_axis], diff --git a/dm_pix/_src/color_conversion.py b/dm_pix/_src/color_conversion.py index 39bc4ac..0c47cf2 100644 --- a/dm_pix/_src/color_conversion.py +++ b/dm_pix/_src/color_conversion.py @@ -43,7 +43,7 @@ def split_channels( chex.assert_axis_dimension(image, axis=channel_axis, expected=3) split_axes = jnp.split(image, 3, axis=channel_axis) - return tuple(map(lambda x: jnp.squeeze(x, axis=channel_axis), split_axes)) + return tuple(map(lambda x: jnp.squeeze(x, axis=channel_axis), split_axes)) # pyrefly: ignore[bad-return] def rgb_to_hsv( @@ -129,9 +129,9 @@ def rgb_planes_to_hsv_planes( norm = 1. / (6. * safe_range) hue = jnp.where(value == green, - norm * (blue - red) + 2. / 6., - norm * (red - green) + 4. / 6.) - hue = jnp.where(value == red, norm * (green - blue), hue) + norm * (blue - red) + 2. / 6., # pyrefly: ignore[unsupported-operation] + norm * (red - green) + 4. / 6.) # pyrefly: ignore[unsupported-operation] + hue = jnp.where(value == red, norm * (green - blue), hue) # pyrefly: ignore[unsupported-operation] hue = jnp.where(range_ > 0, hue, 0.) + (hue < 0.) return hue, saturation, value @@ -159,7 +159,7 @@ def hsv_planes_to_rgb_planes( """ # DO NOT REMOVE - Logging usage. - dh = (hue % 1.0) * 6. # Wrap when hue >= 360°. + dh = (hue % 1.0) * 6. # Wrap when hue >= 360°. # pyrefly: ignore[unsupported-operation] dr = jnp.clip(jnp.abs(dh - 3.) - 1., 0., 1.) dg = jnp.clip(2. - jnp.abs(dh - 2.), 0., 1.) db = jnp.clip(2. - jnp.abs(dh - 4.), 0., 1.) @@ -235,7 +235,7 @@ def hsl_to_rgb( h, s, l = split_channels(image_hsl, channel_axis) - m2 = jnp.where(l <= 0.5, l * (1 + s), l + s - l * s) + m2 = jnp.where(l <= 0.5, l * (1 + s), l + s - l * s) # pyrefly: ignore[unsupported-operation] m1 = 2 * l - m2 def _f(hue): @@ -247,7 +247,7 @@ def _f(hue): jnp.where(hue < 2 / 3, m1 + 6 * (m2 - m1) * (2 / 3 - hue), m1))) image_rgb = jnp.stack([_f(h + 1 / 3), _f(h), _f(h - 1 / 3)], axis=-1) - return jnp.where(s[..., jnp.newaxis] == 0, l[..., jnp.newaxis], image_rgb) + return jnp.where(s[..., jnp.newaxis] == 0, l[..., jnp.newaxis], image_rgb) # pyrefly: ignore[bad-index] def rgb_to_grayscale( diff --git a/dm_pix/_src/color_conversion_test.py b/dm_pix/_src/color_conversion_test.py index 10475d9..4c97583 100644 --- a/dm_pix/_src/color_conversion_test.py +++ b/dm_pix/_src/color_conversion_test.py @@ -110,7 +110,7 @@ def test_rgb_to_hsv(self, test_images, channel_last): for rgb in generate_test_images(*test_images.value): hsv_tf = tf.image.rgb_to_hsv(rgb).numpy() if not channel_last: - rgb = rgb.swapaxes(-1, -3) + rgb = rgb.swapaxes(-1, -3) # pyrefly: ignore[bad-argument-type] hsv_jax = rgb_to_hsv(rgb) if not channel_last: hsv_jax = hsv_jax.swapaxes(-1, -3) @@ -122,7 +122,7 @@ def test_rgb_to_hsv(self, test_images, channel_last): # images to avoid making the tests run too slow. rgb = generate_test_images(*test_images.value)[0] if not channel_last: - rgb = rgb.swapaxes(-1, -3) + rgb = rgb.swapaxes(-1, -3) # pyrefly: ignore[bad-argument-type] jacobian = jax.jacrev(rgb_to_hsv)(rgb) self.assertFalse(jnp.isnan(jacobian).any(), "NaNs in RGB to HSV gradients") @@ -271,7 +271,7 @@ def test_grayscale(self, test_images, keep_dims, channel_last): for rgb in generate_test_images(*test_images.value): grayscale_tf = tf.image.rgb_to_grayscale(rgb).numpy() if not channel_last: - rgb = rgb.swapaxes(-1, -3) + rgb = rgb.swapaxes(-1, -3) # pyrefly: ignore[bad-argument-type] grayscale_jax = rgb_to_grayscale(rgb) if not channel_last: grayscale_jax = grayscale_jax.swapaxes(-1, -3) diff --git a/dm_pix/_src/interpolation.py b/dm_pix/_src/interpolation.py index 4091a4a..9a023d1 100644 --- a/dm_pix/_src/interpolation.py +++ b/dm_pix/_src/interpolation.py @@ -71,15 +71,16 @@ def _make_linear_interpolation_indices_flat_nd( The indices into the flattened input and their weights. """ coordinates = jnp.asarray(coordinates) - shape = jnp.asarray(shape) + shape = jnp.asarray(shape) # pyrefly: ignore[bad-assignment] - if shape.shape[0] != coordinates.shape[0]: + if shape.shape[0] != coordinates.shape[0]: # pyrefly: ignore[missing-attribute] raise ValueError( + # pyrefly: ignore[missing-attribute] (f"{coordinates.shape[0]}-dimensional coordinates provided for " f"{shape.shape[0]}-dimensional input")) lower_nd, upper_nd, weights_nd = _make_linear_interpolation_indices_nd( - coordinates, shape) + coordinates, shape) # pyrefly: ignore[bad-argument-type] # Here we want to translate e.g. a 3D-disposed indices to linear ones, since # we have to index on the flattened source, so: @@ -91,11 +92,11 @@ def _make_linear_interpolation_indices_flat_nd( # column (3rd axis), 2 values to move to get to the same position in the next # row (2nd axis) and 4*2=8 values to move to get to the same position on the # 1st axis. - strides = jnp.concatenate([jnp.cumprod(shape[:0:-1])[::-1], jnp.array([1])]) + strides = jnp.concatenate([jnp.cumprod(shape[:0:-1])[::-1], jnp.array([1])]) # pyrefly: ignore[bad-argument-type] # Array of 2^n rows where the ith row is the binary representation of i. binary_array = jnp.array( - list(itertools.product([0, 1], repeat=shape.shape[0]))) + list(itertools.product([0, 1], repeat=shape.shape[0]))) # pyrefly: ignore[missing-attribute] # Expand dimensions to allow broadcasting `strides` and `binary_array` to # every coordinate. @@ -127,7 +128,7 @@ def _linear_interpolate_using_indices_nd( weights: chex.Array, ) -> chex.Array: """Interpolates linearly on `volume` using `indices` and `weights`.""" - target = jnp.sum(weights * volume[indices], axis=0) + target = jnp.sum(weights * volume[indices], axis=0) # pyrefly: ignore[bad-index] if jnp.issubdtype(volume.dtype, jnp.integer): target = _round_half_away_from_zero(target) return target.astype(volume.dtype) @@ -205,12 +206,12 @@ def flat_nd_linear_interpolate_constant( # boundaries. is_in_bounds = jnp.full(coordinates.shape[1:], True) for dim, dim_size in enumerate(volume_shape): - is_in_bounds = jnp.logical_and(is_in_bounds, coordinates[dim] >= 0) + is_in_bounds = jnp.logical_and(is_in_bounds, coordinates[dim] >= 0) # pyrefly: ignore[bad-index] is_in_bounds = jnp.logical_and(is_in_bounds, - coordinates[dim] <= dim_size - 1) + coordinates[dim] <= dim_size - 1) # pyrefly: ignore[bad-index] return flat_nd_linear_interpolate( volume, coordinates, unflattened_vol_shape=unflattened_vol_shape - ) * is_in_bounds + (1. - is_in_bounds) * cval + ) * is_in_bounds + (1. - is_in_bounds) * cval # pyrefly: ignore[unsupported-operation] diff --git a/dm_pix/_src/metrics.py b/dm_pix/_src/metrics.py index 4468728..6dfd6be 100644 --- a/dm_pix/_src/metrics.py +++ b/dm_pix/_src/metrics.py @@ -48,7 +48,7 @@ def mae( chex.assert_type([a, b], float) chex.assert_equal_shape([a, b]) mean_fn = jnp.nanmean if ignore_nans else jnp.mean - return mean_fn(jnp.abs(a - b), axis=(-3, -2, -1)) + return mean_fn(jnp.abs(a - b), axis=(-3, -2, -1)) # pyrefly: ignore[unsupported-operation] def mse( @@ -73,7 +73,7 @@ def mse( chex.assert_type([a, b], float) chex.assert_equal_shape([a, b]) mean_fn = jnp.nanmean if ignore_nans else jnp.mean - return mean_fn(jnp.square(a - b), axis=(-3, -2, -1)) + return mean_fn(jnp.square(a - b), axis=(-3, -2, -1)) # pyrefly: ignore[unsupported-operation] def psnr( @@ -269,9 +269,9 @@ def filter_fn_x(z): mu00 = mu0 * mu0 mu11 = mu1 * mu1 mu01 = mu0 * mu1 - sigma00 = filter_fn(a**2) - mu00 - sigma11 = filter_fn(b**2) - mu11 - sigma01 = filter_fn(a * b) - mu01 + sigma00 = filter_fn(a**2) - mu00 # pyrefly: ignore[unsupported-operation] + sigma11 = filter_fn(b**2) - mu11 # pyrefly: ignore[unsupported-operation] + sigma01 = filter_fn(a * b) - mu01 # pyrefly: ignore[unsupported-operation] # Clip the variances and covariances to valid values. # Variance must be non-negative: diff --git a/examples/image_augmentation.py b/examples/image_augmentation.py index c3024d0..e881a32 100644 --- a/examples/image_augmentation.py +++ b/examples/image_augmentation.py @@ -31,7 +31,7 @@ def main(_) -> None: # Flip up-down the image and visual it. flip_up_down_image = pix.flip_up_down(image=image) - _imshow(flip_up_down_image) + _imshow(flip_up_down_image) # pyrefly: ignore[bad-argument-type] # Apply a Gaussian filter to the image and visual it. gaussian_blur_image = pix.gaussian_blur( @@ -39,35 +39,35 @@ def main(_) -> None: sigma=_KERNEL_SIGMA, kernel_size=_KERNEL_SIZE, ) - _imshow(gaussian_blur_image) + _imshow(gaussian_blur_image) # pyrefly: ignore[bad-argument-type] # Change image brightness and visual it. adjust_brightness_image = pix.adjust_brightness( image=image, delta=_MAGIC_VALUE, ) - _imshow(adjust_brightness_image) + _imshow(adjust_brightness_image) # pyrefly: ignore[bad-argument-type] # Change image contrast and visual it. adjust_contrast_image = pix.adjust_contrast( image=image, factor=_MAGIC_VALUE, ) - _imshow(adjust_contrast_image) + _imshow(adjust_contrast_image) # pyrefly: ignore[bad-argument-type] # Change image gamma and visual it. adjust_gamma_image = pix.adjust_gamma( image=image, gamma=_MAGIC_VALUE, ) - _imshow(adjust_gamma_image) + _imshow(adjust_gamma_image) # pyrefly: ignore[bad-argument-type] # Change image hue and visual it. adjust_hue_image = pix.adjust_hue( image=image, delta=_MAGIC_VALUE, ) - _imshow(adjust_hue_image) + _imshow(adjust_hue_image) # pyrefly: ignore[bad-argument-type] def _get_image():