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
Fixes 500+ mislabeled MURA images

In this repository, new csv files are provided that fixes 500+ mislabeled MURA x-rays for all categories. The mislabeled x-rays mainly had hardware in them. This project only fixes the false negative

Pieter Zeilstra 4 May 18, 2022
missing-pixel-filler is a python package that, given images that may contain missing data regions (like satellite imagery with swath gaps), returns these images with the regions filled.

Missing Pixel Filler This is the official code repository for the Missing Pixel Filler by SpaceML. missing-pixel-filler is a python package that, give

SpaceML 11 Jul 19, 2022
🛹 Turn an SVG into an STL for stencil creation purposes

svg2stl This repository provides a script which takes as input an SVG such as this one: It outputs an STL file like this one: You can also see an inte

Max Halford 3 Dec 29, 2021
A tool for making simple-style text posters or wallpapers with high resolution.

PurePoster PurePoster is a fancy tool for making arbitrary-resolution, simple-style posters or wallpapers with text in center. Functionality PurePoste

Renyang Guan 4 Jul 09, 2022
Panel Competition Image Generator

Panel Competition Image Generator This project was build by a member of the NFH community and is open for everyone who wants to try it. Relevant links

Juliano Mendieta 1 Oct 22, 2021
3D Reconstruction Software

Meshroom is a free, open-source 3D Reconstruction Software based on the AliceVision Photogrammetric Computer Vision framework. Learn more details abou

AliceVision 8.7k Jan 02, 2023
Tool made for the FWA Yearbook Team to resize multiple images quickly.

ImageResize Tool Tool made for the FWA Yearbook Team to resize multiple images quickly. Make sure to check this repo for future updates How to Use The

LGobin 1 Jan 07, 2022
Turtle graphics || Python

turtle Turtle graphics || Python Rainbow (রংধনু) : Rainbow.using.Python.--.Python.Turtle.graphics.mp4 Human robot (মানব রোবট) : Draw.a.human.robot.usi

Jubair Ahmed Junjun 1 Oct 08, 2021
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
A GUI-based (PyQt5) tool used to design 2D linkage mechanism.

Pyslvs-UI A GUI-based (PyQt5) tool used to design 2D linkage mechanism. Planar Linkages Simulation Python-Solvespace: Kernel from Solvespace with Cyth

Yuan Chang 141 Dec 13, 2022
Detecting haze image with hazer.

hazer-py Detecting haze image with hazer. What is hazer Hazer is a lib for getting "haze degree". This repository is python version of hazer: https://

Joey777210 2 Dec 27, 2021
Program for analyzing shadows from Cassini images

Moons: An Analysis Module for Vicar Files General This packages/program was created for my bachelor's thesis for the Astronomy department at Universit

Joni 1 Jul 16, 2021
Fuzzware is a project for automated, self-configuring fuzzing of firmware images

Fuzzware Fuzzware is a project for automated, self-configuring fuzzing of firmware images. The idea of this project is to configure the memory ranges

190 Dec 21, 2022
Bot by image recognition simulating (random) human clicks

bbbot22 bot por reconhecimento de imagem simulando cliques humanos (aleatórios) inb4: sim, esse é basicamente o mesmo bot de 2021 porque a Globo não t

Yuri 2 Apr 05, 2022
HTML2Image is a lightweight Python package that acts as a wrapper around the headless mode of existing web browsers to generate images from URLs and from HTML+CSS strings or files.

A package acting as a wrapper around the headless mode of existing web browsers to generate images from URLs and from HTML+CSS strings or files.

176 Jan 01, 2023
QR-code Generator with a basic GUI.

Qr_generator_python Qr code generator with a basic GUI. ❔ About the QR-Code-Generator This project Generates QR codes to sites, e-mails and plain text

Tecixck 2 Oct 11, 2021
Black-white image converter - Black-white photo colorization

Black-white image converter - Black-white photo colorization

1 Jan 02, 2022
A sketch like(?) effect for images

lineArt A sketch like(?) effect for images How to run main.py [filename] [option {1,2}] option 1 retains colour option 2 gives gray image #results ori

1 Oct 28, 2021
With this simple py script you will be able to get all the .png from a folder and generate a yml for Oraxen

Oraxen-item-to-yml With this simple py script you will be able to get all the .png from a folder and generate a yml for Oraxen How to use Install the

Akex 1 Dec 29, 2021
Generate your own QR Code and scan it to see the results! Never use it for malicious purposes.

QR-Code-Generator-Python Choose the name of your generated QR .png file. If it happens to open the .py file (the application), there are specific comm

1 Dec 23, 2021