PIX is an image processing library in JAX, for JAX.

Overview

PIX

PIX is an image processing library in JAX, for JAX.

Overview

JAX is a library resulting from the union of Autograd and XLA for high-performance machine learning research. It provides NumPy, SciPy, automatic differentiation and first-class GPU/TPU support.

PIX is a library built on top of JAX with the goal of providing image processing functions and tools to JAX in a way that they can be optimised and parallelised through jax.jit, jax.vmap and jax.pmap.

Installation

PIX is written in pure Python, but depends on C++ code via JAX.

Because JAX installation is different depending on your CUDA version, PIX does not list JAX as a dependency in requirements.txt, although it is technically listed for reference, but commented.

First, follow JAX installation instructions to install JAX with the relevant accelerator support.

Then, install PIX using pip:

$ pip install git+https://github.com/deepmind/dm_pix

Quickstart

To use PIX, you just need to import dm_pix as pix and use it right away!

For example, let's assume to have loaded the JAX logo (available in examples/assets/jax_logo.jpg) in a variable called image and we want to flip it left to right.

JAX logo

All it's needed is the following code!

import dm_pix as pix

# Load an image into a NumPy array with your preferred library.
image = load_image()

flip_left_right_image = pix.flip_left_right(image)

And here is the result!

JAX logo left-right

All the functions in PIX can be jax.jited, jax.vmaped and jax.pmaped, so all the following functions can take advantage of optimization and parallelization.

import dm_pix as pix
import jax

# Load an image into a NumPy array with your preferred library.
image = load_image()

# Vanilla Python function.
flip_left_right_image = pix.flip_left_right(image)

# `jax.jit`ed function.
flip_left_right_image = jax.jit(pix.flip_left_right)(image)

# Assuming to have a single device, like a CPU or a single GPU, we add a
# single leading dimension for using `image` with the parallelized or
# the multi-device parallelization version of `pix.flip_left_right`.
# To know more, please refer to JAX documentation of `jax.vmap` and `jax.pmap`.
image = image[np.newaxis, ...]

# `jax.vmap`ed function.
flip_left_right_image = jax.vmap(pix.flip_left_right)(image)

# `jax.pmap`ed function.
flip_left_right_image = jax.pmap(pix.flip_left_right)(image)

You can check it yourself that the result from the four versions of pix.flip_left_right is the same (up to the accelerator floating point accuracy)!

Examples

We have a few examples in the examples/ folder. They are not much more involved then the previous example, but they may be a good starting point for you!

Testing

We provide a suite of tests to help you both testing your development environment and to know more about the library itself! All test files have _test suffix, and can be executed using pytest.

If you already have PIX installed, you just need to install some extra dependencies and run pytest as follows:

$ pip install -r requirements_tests.txt
$ python -m pytest [-n <NUMCPUS>] dm_pix

If you want an isolated virtual environment, you just need to run our utility bash script as follows:

$ ./test.sh

Citing PIX

This repository is part of the DeepMind JAX Ecosystem, to cite PIX please use the DeepMind JAX Ecosystem citation.

Contribute!

We are very happy to accept contributions!

Please read our contributing guidelines and send us PRs!

Comments
  • Affine transform of RGB image

    Affine transform of RGB image

    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: image

    documentation question 
    opened by jloganolson 7
  • Add probability param for random flip functions.

    Add probability param for random flip functions.

    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.

    opened by pabloduque0 5
  • Transformations for 3d images

    Transformations for 3d images

    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)?
    opened by jakubMitura14 3
  • About random gamma

    About random gamma

    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?

    question 
    opened by alonfnt 3
  • Import error when upgrading to jax 0.3.18

    Import error when upgrading to jax 0.3.18

    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'
    
    opened by ASEM000 2
  • Add random_gamma to the augmentation functions

    Add random_gamma to the augmentation functions

    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.

    opened by alonfnt 2
  • Image resizing operation

    Image resizing operation

    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 ?

    opened by jayendra13 2
  • Fix NaNs appearing in gradients through HSV conversion.

    Fix NaNs appearing in gradients through HSV conversion.

    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.

    opened by copybara-service[bot] 1
  • Internal changes.

    Internal changes.

    Internal changes.

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

    opened by copybara-service[bot] 1
  • Bump pillow from 9.0.0 to 9.0.1

    Bump pillow from 9.0.0 to 9.0.1

    Bumps pillow from 9.0.0 to 9.0.1.

    Release notes

    Sourced from pillow's releases.

    9.0.1

    https://pillow.readthedocs.io/en/stable/releasenotes/9.0.1.html

    Changes

    • In show_file, use os.remove to remove temporary images. CVE-2022-24303 #6010 [@​radarhere, @​hugovk]
    • Restrict builtins within lambdas for ImageMath.eval. CVE-2022-22817 #6009 [radarhere]
    Changelog

    Sourced from pillow's changelog.

    9.0.1 (2022-02-03)

    • In show_file, use os.remove to remove temporary images. CVE-2022-24303 #6010 [radarhere, hugovk]

    • Restrict builtins within lambdas for ImageMath.eval. CVE-2022-22817 #6009 [radarhere]

    Commits
    • 6deac9e 9.0.1 version bump
    • c04d812 Update CHANGES.rst [ci skip]
    • 4fabec3 Added release notes for 9.0.1
    • 02affaa Added delay after opening image with xdg-open
    • ca0b585 Updated formatting
    • 427221e In show_file, use os.remove to remove temporary images
    • c930be0 Restrict builtins within lambdas for ImageMath.eval
    • 75b69dd Dont need to pin for GHA
    • cd938a7 Autolink CWE numbers with sphinx-issues
    • 2e9c461 Add CVE IDs
    • See full diff in compare view

    Dependabot compatibility score

    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.


    Dependabot commands and options

    You can trigger Dependabot actions by commenting on this PR:

    • @dependabot rebase will rebase this PR
    • @dependabot recreate will recreate this PR, overwriting any edits that have been made to it
    • @dependabot merge will merge this PR after your CI passes on it
    • @dependabot squash and merge will squash and merge this PR after your CI passes on it
    • @dependabot cancel merge will cancel a previously requested merge and block automerging
    • @dependabot reopen will reopen this PR if it is closed
    • @dependabot close will close this PR and stop Dependabot recreating it. You can achieve the same result by closing it manually
    • @dependabot ignore this major version will close this PR and stop Dependabot creating any more for this major version (unless you reopen the PR or upgrade to it yourself)
    • @dependabot ignore this minor version will close this PR and stop Dependabot creating any more for this minor version (unless you reopen the PR or upgrade to it yourself)
    • @dependabot ignore this dependency will close this PR and stop Dependabot creating any more for this dependency (unless you reopen the PR or upgrade to it yourself)
    • @dependabot use these labels will set the current labels as the default for future PRs for this repo and language
    • @dependabot use these reviewers will set the current reviewers as the default for future PRs for this repo and language
    • @dependabot use these assignees will set the current assignees as the default for future PRs for this repo and language
    • @dependabot use this milestone will set the current milestone as the default for future PRs for this repo and language

    You can disable automated security fix PRs for this repo from the Security Alerts page.

    dependencies 
    opened by dependabot[bot] 0
  • Bump ipython from 7.16.1 to 7.16.3

    Bump ipython from 7.16.1 to 7.16.3

    Bumps ipython from 7.16.1 to 7.16.3.

    Commits

    Dependabot compatibility score

    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.


    Dependabot commands and options

    You can trigger Dependabot actions by commenting on this PR:

    • @dependabot rebase will rebase this PR
    • @dependabot recreate will recreate this PR, overwriting any edits that have been made to it
    • @dependabot merge will merge this PR after your CI passes on it
    • @dependabot squash and merge will squash and merge this PR after your CI passes on it
    • @dependabot cancel merge will cancel a previously requested merge and block automerging
    • @dependabot reopen will reopen this PR if it is closed
    • @dependabot close will close this PR and stop Dependabot recreating it. You can achieve the same result by closing it manually
    • @dependabot ignore this major version will close this PR and stop Dependabot creating any more for this major version (unless you reopen the PR or upgrade to it yourself)
    • @dependabot ignore this minor version will close this PR and stop Dependabot creating any more for this minor version (unless you reopen the PR or upgrade to it yourself)
    • @dependabot ignore this dependency will close this PR and stop Dependabot creating any more for this dependency (unless you reopen the PR or upgrade to it yourself)
    • @dependabot use these labels will set the current labels as the default for future PRs for this repo and language
    • @dependabot use these reviewers will set the current reviewers as the default for future PRs for this repo and language
    • @dependabot use these assignees will set the current assignees as the default for future PRs for this repo and language
    • @dependabot use this milestone will set the current milestone as the default for future PRs for this repo and language

    You can disable automated security fix PRs for this repo from the Security Alerts page.

    dependencies 
    opened by dependabot[bot] 0
  • Gaussian blur cpu performance

    Gaussian blur cpu performance

    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.

    enhancement 
    opened by mgoulao 12
Releases(v0.3.4)
  • v0.3.4(Sep 12, 2022)

    What's Changed

    • Implement interpolation with constant value and rotation by an arbitrary angle, by @swagnercarena in https://github.com/deepmind/dm_pix/pull/42
    • Exposed interpolation with constant value and rotation by arbitrary angles when importing dm-pix, by @swagnercarena in https://github.com/deepmind/dm_pix/pull/44
    • Add elastic deformation, by @pabloduque0 in https://github.com/deepmind/dm_pix/pull/45

    New Contributors

    • @swagnercarena made their first contribution in https://github.com/deepmind/dm_pix/pull/42

    Full Changelog: https://github.com/deepmind/dm_pix/compare/v0.3.3...v0.3.4

    Source code(tar.gz)
    Source code(zip)
  • v0.3.3(Aug 1, 2022)

    What's Changed

    • Add probability param for random flip functions, by @pabloduque0 in https://github.com/deepmind/dm_pix/pull/38
    • Add general affine transform, by @copybara-service in https://github.com/deepmind/dm_pix/pull/39
    • Bump version to 0.3.3 by @copybara-service in https://github.com/deepmind/dm_pix/pull/40

    New Contributors

    • @pabloduque0 made their first contribution in https://github.com/deepmind/dm_pix/pull/38

    Full Changelog: https://github.com/deepmind/dm_pix/compare/v0.3.2...v0.3.3

    Source code(tar.gz)
    Source code(zip)
  • v0.3.2(May 5, 2022)

    What's Changed

    • Fix ReadTheDocs build by @copybara-service in https://github.com/deepmind/dm_pix/pull/35

    Full Changelog: https://github.com/deepmind/dm_pix/compare/v0.3.1...v0.3.2

    Source code(tar.gz)
    Source code(zip)
  • v0.3.1(May 4, 2022)

    What's Changed

    • Add ReadTheDocs documentation by @copybara-service in https://github.com/deepmind/dm_pix/pull/22
    • Fix ReadTheDocs documentation by @copybara-service in https://github.com/deepmind/dm_pix/pull/23
    • Update README with GitHub shields by @copybara-service in https://github.com/deepmind/dm_pix/pull/24
    • Update README shields by @copybara-service in https://github.com/deepmind/dm_pix/pull/25
    • Added interactive examples of PIX Augmentations with Google Colab by @SupreethRao99 in https://github.com/deepmind/dm_pix/pull/26
    • Bump pillow from 8.3.2 to 9.0.0 by @dependabot in https://github.com/deepmind/dm_pix/pull/29
    • Migrate away from using JaxTestCase in tests by @copybara-service in https://github.com/deepmind/dm_pix/pull/30
    • Fix pix's SSIM's gradient to not nan-out on a flat image, and add a unit test that catches it. by @copybara-service in https://github.com/deepmind/dm_pix/pull/31
    • Bump ipython from 7.16.1 to 7.16.3 by @dependabot in https://github.com/deepmind/dm_pix/pull/32
    • Bump pillow from 9.0.0 to 9.0.1 by @dependabot in https://github.com/deepmind/dm_pix/pull/33
    • Bump version to 0.3.1 by @copybara-service in https://github.com/deepmind/dm_pix/pull/34

    New Contributors

    • @SupreethRao99 made their first contribution in https://github.com/deepmind/dm_pix/pull/26

    Full Changelog: https://github.com/deepmind/dm_pix/compare/v0.3.0...v0.3.1

    Source code(tar.gz)
    Source code(zip)
  • v0.3.0(Nov 18, 2021)

  • v0.2.1(Nov 18, 2021)

  • v0.2.0(Nov 1, 2021)

    Changelog

    Full Changelog

    Changes

    * This Changelog was automatically generated by github_changelog_generator and manually updated by the project devs.

    Source code(tar.gz)
    Source code(zip)
  • v0.1.0(Jul 15, 2021)

Owner
DeepMind
DeepMind
impy is an all-in-one image analysis library, equipped with parallel processing, GPU support, GUI based tools and so on.

impy is All You Need in Image Analysis impy is an all-in-one image analysis library, equipped with parallel processing, GPU support, GUI based tools a

null 24 Dec 20, 2022
starfish is a Python library for processing images of image-based spatial transcriptomics.

starfish: scalable pipelines for image-based transcriptomics starfish is a Python library for processing images of image-based spatial transcriptomics

null 199 Dec 8, 2022
Blue noise image stippling in Processing (requires GPU)

Blue noise image stippling in Processing (requires GPU)

Mike Wong 141 Oct 9, 2022
A Python package implementing various HDRI / Radiance image processing algorithms.

Colour - HDRI A Python package implementing various HDRI / Radiance image processing algorithms. It is open source and freely available under the New

colour-science 111 Dec 6, 2022
GPU-accelerated image processing using cupy and CUDA

napari-cupy-image-processing GPU-accelerated image processing using cupy and CUDA This napari plugin was generated with Cookiecutter using with @napar

Robert Haase 16 Oct 26, 2022
image-processing exercises.

image_processing Assignment 21 Checkered Board Create a chess table using numpy and opencv. view: Color Correction Reverse black and white colors with

Benyamin Zojaji 25 Dec 15, 2022
clesperanto is a graphical user interface for GPU-accelerated image processing.

clesperanto is a graphical user interface for a multi-platform multi-language framework for GPU-accelerated image processing. It is based on napari and the pyclesperanto-prototype.

null 1 Jan 2, 2022
Simple Python image processing & automatization project for a simple web based game

What is this? Simple Python image processing & automatization project for a simple web based game Made using only Github Copilot (except the color and

SGeri 2 Aug 15, 2022
Image Processing - Make noise images clean

影像處理-影像降躁化(去躁化) (Image Processing - Make Noise Images Clean) 得力於電腦效能的大幅提升以及GPU的平行運算架構,讓我們能夠更快速且有效地訓練AI,並將AI技術應用於不同領域。本篇將帶給大家的是 「將深度學習應用於影像處理中的影像降躁化 」,

null 2 Aug 4, 2022
Napari simpleitk image processing

napari-simpleitk-image-processing (n-SimpleITK) Process images using SimpleITK in napari Usage Filters of this napari plugin can be found in the Tools

Robert Haase 11 Dec 19, 2022
A small Python module for BMP image processing.

micropython-microbmp A small Python module for BMP image processing. It supports BMP image of 1/2/4/8/24-bit colour depth. Loading supports compressio

Quan Lin 4 Nov 2, 2022
Image Processing HighPass Filter With Python

Image_Processing_HighPassFilter High Pass Filter take the high frequency and ignore the low frequency High Pass Filter can be use to sharpening an ima

Felix Pratamasan 1 Dec 27, 2021
kikuchipy is an open-source Python library for processing and analysis of electron backscatter diffraction (EBSD) patterns

kikuchipy is an open-source Python library for processing and analysis of electron backscatter diffraction (EBSD) patterns. The library builds on the

pyxem 53 Dec 29, 2022
Seaborn-image is a Python image visualization library based on matplotlib and provides a high-level API to draw attractive and informative images quickly and effectively.

seaborn-image: image data visualization Description Seaborn-image is a Python image visualization library based on matplotlib and provides a high-leve

null 48 Jan 5, 2023
Using P5.js, Processing and Python to create generative art

Experiments in Generative Art Using Python, Processing, and P5.js Quick Links Daily Sketches March 2021. | Gallery | Repo | Done using P5.js Genuary 2

Ram Narasimhan 33 Jul 6, 2022
Pixel Brush Processing Unit

Pixel Brush Processing Unit The Pixel Brush Processing Unit (PBPU for short) is a simple 4-Bit CPU I designed in Logisim while I was still in school a

Pixel Brush 2 Nov 3, 2022
An executor that performs standard pre-processing and normalization on images.

An executor that performs standard pre-processing and normalization on images.

Jina AI 6 Jun 30, 2022
An open source image editor which can manipulate an image in many ways!

Image Editor - An open source image editor which can manipulate an image in many ways! If you need any more modes in repo or I

TroJanzHEX 44 Nov 17, 2022
Image enhancing model for making a blurred image to be somehow clearer than before

This is a very small prject which helps in enhancing the images by taking a Input images. This project has many features like detcting the faces and enhaning the faces itself and also a feature which enhances the whole image

null 3 Dec 3, 2021