Hello! Do you have any idea why this basic translation matrix is affecting each channel of an RGB image differently?

img = jp.array(PIL.Image.open('test.png').convert('RGB')) / 255.
print(img.shape)
tx = 100.
ty = 0.
translate_matrix = jp.array([
    [1,0,ty],
    [0,1,tx],
    [0,0,1],
])

t_img = pix.affine_transform(img,translate_matrix,mode="constant")
plt.imshow(t_img)

Screenshot of output:

Based on the conversation in #37 . Will work on adding one or two functions in the upcoming weeks.

Apologies for the extra unnecessary commits, vs studio butcher the notebook and had to revert that one, we can squash it. Happy to update the notebook in a follow up as well.

A few thoughts:

• I have not added checks (asserts) on the prob value to allow jitting without needing static args. Let me know if this is okay, Im not super familiar with chex so maybe Im missing something there.
• I had to change the way the extra parameters are evaluated in the test to kwargs always. I think this should be fine as long as we dont make some of those parameter positional only. Let me know if you rather take a different approach on this one.
• I love how clever the test setup is! 🚀🚀🚀

Feel free to add any comments on the wording of the docstring as well.

Hello 2 questions

1. as the transformations are implemented in pure Jax they are differentiable in most cases - am I right? (more precisely I am considering putting the probability of applying transform as learnable parameter)
2. is it possible to use transformations in 3D images (Magnetic resonance imagiing)?

Is there a reason adjust_gamma is missing its random equivalent as in the case of adjust_contrast, adjust_brightness, adjust_hue and adjust_saturation?

Hi,

jax version = 0.3.18 I am faced with the following error when importing dm_pix

I suspect it's due to the removal of ._src since jax 0.3.18 https://github.com/google/jax/releases/tag/jax-v0.3.18

----> [1](vscode-notebook-cell:/Users/asem/serket/test.ipynb#Y121sZmlsZQ%3D%3D?line=0) import dm_pix

File ~/miniforge3/envs/dev-jax15/lib/python3.10/site-packages/dm_pix/__init__.py:16, in <module>
      1 # Copyright 2020 DeepMind Technologies Limited. All Rights Reserved.
      2 #
      3 # Licensed under the Apache License, Version 2.0 (the "License");
   (...)
     12 # See the License for the specific language governing permissions and
     13 # limitations under the License.
     14 """PIX public APIs."""
---> 16 from dm_pix._src import augment
     17 from dm_pix._src import color_conversion
     18 from dm_pix._src import depth_and_space

File ~/miniforge3/envs/dev-jax15/lib/python3.10/site-packages/dm_pix/_src/augment.py:25, in <module>
     22 import functools
     23 from typing import Callable, Sequence, Tuple, Union
---> 25 import chex
     26 from dm_pix._src import color_conversion
     27 from dm_pix._src import interpolation

File ~/miniforge3/envs/dev-jax15/lib/python3.10/site-packages/chex/__init__.py:17, in <module>
      1 # Copyright 2020 DeepMind Technologies Limited. All Rights Reserved.
      2 #
      3 # Licensed under the Apache License, Version 2.0 (the "License");
   (...)
     13 # limitations under the License.
     14 # ==============================================================================
     15 """Chex: Testing made fun, in JAX!"""
---> 17 from chex._src.asserts import assert_axis_dimension
     18 from chex._src.asserts import assert_axis_dimension_comparator
     19 from chex._src.asserts import assert_axis_dimension_gt

File ~/miniforge3/envs/dev-jax15/lib/python3.10/site-packages/chex/_src/asserts.py:26, in <module>
     23 import unittest
     24 from unittest import mock
---> 26 from chex._src import asserts_internal as _ai
     27 from chex._src import pytypes
     28 import jax

File ~/miniforge3/envs/dev-jax15/lib/python3.10/site-packages/chex/_src/asserts_internal.py:32, in <module>
     29 from typing import Any, Sequence, Union, Callable, Optional, Set, Tuple, Type
     31 from absl import logging
---> 32 from chex._src import pytypes
     33 import jax
     34 import jax.numpy as jnp

File ~/miniforge3/envs/dev-jax15/lib/python3.10/site-packages/chex/_src/pytypes.py:44, in <module>
     40 Device = jax.lib.xla_extension.Device
     42 ArrayTree = Union[Array, Iterable['ArrayTree'], Mapping[Any, 'ArrayTree']]
---> 44 ArrayDType = jax._src.numpy.lax_numpy._ScalarMeta

AttributeError: module 'jax' has no attribute '_src'

Since I have several projects where I have used this wrapper, and following on #43, I thought it would be okay to submit it as a PR to have it already in PIX.

random_gamma is just a wrapper for adjust_gamma where the value of the gamma parameter is sampled uniformly in the given range, similarly to what other augmentation functions already do in the library.

Is there any specific reason not to include the resize operation ?

https://www.tensorflow.org/api_docs/python/tf/image/resize

There is a interpolation module in the code and would it be useful for this tf.image.resize operation ?

Fix NaNs appearing in gradients through HSV conversion.

When range_ or value are 0., division by zeros appear, causing NaNs. This is not a problem when computing the function because these are masked by jnp.where, but this causes problems during gradient computation.

Also included some test to highlight the problems before the fix. The jacobian is only evaluated for the last image of all the TestImages, to avoid slowing down the tests to much. ALL_ONES and ALL_ZEROS highlight the problem.

Internal changes.

FUTURE_COPYBARA_INTEGRATE_REVIEW=https://github.com/deepmind/dm_pix/pull/26 from SupreethRao99:master a060f0c023539b0fc8cad0ff0ab6fbd6cc7289ca

Bumps pillow from 9.0.0 to 9.0.1.

Dependabot will resolve any conflicts with this PR as long as you don't alter it yourself. You can also trigger a rebase manually by commenting @dependabot rebase.

Bumps ipython from 7.16.1 to 7.16.3.

Dependabot will resolve any conflicts with this PR as long as you don't alter it yourself. You can also trigger a rebase manually by commenting @dependabot rebase.

I have been doing some experiments with PIX since it allows computing image augmentations in the GPU in contrast to torchvision which computes in the CPU and requires multiple workers to avoid bottlenecks. When performing some very simple timeit examples I observed a very high time when performing a gaussian blur in the CPU. I created a simple Colab notebook to demonstrate these experiments. I even tested transferring the image to CPU before performing the blur but it doesn't seem to make any difference. I was wondering if this is intended and I should not rely on CPU computations at all or if something is yet to be optimized for CPU computation.