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
Console images in 48 colors, 216 colors and full rgb

console_images Console images in 48 colors, 216 colors and full rgb Full RGB 216 colors 48 colors If it does not work maybe you should change color_fu

Урядов Алексей 5 Oct 11, 2022
A suite of useful tools based on 3D interactivity in napari

napari-threedee A suite of useful tools based on 3D interactivity in napari This napari plugin was generated with Cookiecutter using @napari's cookiec

11 Dec 14, 2022
Labelme is a graphical image annotation tool, It is written in Python and uses Qt for its graphical interface

Image Polygonal Annotation with Python (polygon, rectangle, circle, line, point and image-level flag annotation).

Kentaro Wada 9.6k Jan 09, 2023
Archive of the image generator stuff from my API

alex_api_archive Archive of the image generator stuff from my API FAQ Q: Why? A: Because I am removing these components from the API Q: How do I run i

AlexFlipnote 26 Nov 17, 2022
The friendly PIL fork (Python Imaging Library)

Pillow Python Imaging Library (Fork) Pillow is the friendly PIL fork by Alex Clark and Contributors. PIL is the Python Imaging Library by Fredrik Lund

Pillow 10.4k Dec 31, 2022
Python scripts for semi-automated morphometric analysis of atolls from Landsat satellite Imagery.

AtollGeoMorph Python scripts for semi-automated morphometric analysis of atolls from Landsat satellite Imagery. The python scripts included allow user

1 Dec 16, 2022
This is a python project which detects color of an image when you double click on it.

This is a python project which detects color of an image when you double click on it. You have to press ESC button to close the pop-up Image window. There are mainly two library CV2 and Pandas that a

Yashwant Kumar Singh 0 Aug 16, 2022
Python-fu-cartoonify - GIMP plug-in to turn a photo into a cartoon.

python-fu-cartoonify GIMP plug-in to turn a photo into a cartoon. Preview Installation Copy python-fu-cartoonify.py into the plug-in folder listed und

Pascal Reitermann 6 Aug 05, 2022
A Icon Maker GUI Made - Convert your image into icon ( .ico format ).

Icon-Maker-GUI A Icon Maker GUI Made Using Python 3.9.0 . It will take any image and convert it to ICO file, for web site favicon or Windows applicati

Insanecodes 12 Dec 15, 2021
Validate arbitrary image uploads from incoming data urls while preserving file integrity but removing EXIF and unwanted artifacts and RCE exploit potential

Validate arbitrary base64-encoded image uploads as incoming data urls while preserving image integrity but removing EXIF and unwanted artifacts and mitigating RCE-exploit potential.

A3R0 1 Jan 10, 2022
Plots the graph of a function with ASCII characters.

ASCII Graph Plotter Plots the graph of a function with ASCII characters. See the change log here. Developed by InformaticFreak (c) 2021 How to use py

InformaticFreak 2 Apr 29, 2022
FrostedGlass is a translucent frosted glass effect widget, that creates a context with the background behind it.

FrostedGlass FrostedGlass is a translucent frosted glass effect widget, that creates a context with the background behind it. The effect is drawn on t

Kivy Garden 24 Dec 10, 2022
Simple utility to tinker with OPlus images

OPlus image utilities Prerequisites Linux running kernel 5.4 or up (check with uname -r) Image rebuilding Used to rebuild read-only erofs images into

Wiley Lau 15 Dec 28, 2022
Viewer for NFO files

NFO Viewer NFO Viewer is a simple viewer for NFO files, which are "ASCII" art in the CP437 codepage. The advantages of using NFO Viewer instead of a t

Osmo Salomaa 114 Dec 29, 2022
Simple AI app that is guessing color of apple in picture

Apple Color Determinant Application that is guessing color of apple from image Install Pillow, sklearn and numpy, using command for your package manag

Gleb Nikitin 1 Oct 25, 2021
Depix is a tool for recovering passwords from pixelized screenshots.

This implementation works on pixelized images that were created with a linear box filter. In this article I cover background information on pixelization and similar research.

23.1k Jan 04, 2023
Alternate Python bindings for the Open Asset Import Library (ASSIMP)

Impasse A simple Python wrapper for assimp using cffi to access the library. Requires Python = 3.7. It's a fork of PyAssimp, Assimp's official Python

Salad Dais 3 Sep 26, 2022
Simple program to easily view Euler parameters in 3D.

Simple program to easily view Euler parameters in 3D.

5 Aug 20, 2021
将位图转为彩色矢量 svg 图片

一个将位图描摹为彩色矢量 svg 图片的程序,是一个命令行工具,使用 Python 脚本实现,运行环境 Python3.8+。 ✨ 效果 以一个字帖图片为例,这是 png 格式的位图(370KB): 这是颜

Haujet Zhao 104 Dec 30, 2022
This will help to read QR codes using Raspberry Pi and Pi Camera

Raspberry-Pi-Generate-and-Read-QR-code This will help to read QR codes using Raspberry Pi and Pi Camera Install the required libraries first in your T

Raspberry_Pi Pakistan 2 Nov 06, 2021