google-deepmind/dm_pix

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

What is PIX?

PIX is a lightweight image‑processing library built on top of JAX, the high‑performance NumPy‑compatible framework that powers a lot of modern machine‑learning research. All of PIX’s functions (e.g., flipping, resizing, color conversion) are written so they can be compiled with jax.jit, vectorised with jax.vmap, or run on multiple devices with jax.pmap. In practice this means you can treat images just like any other JAX array and get the same GPU/TPU speed‑ups you get for neural‑network code.

Quick start

import dm_pix as pix

# Load an image with any library that returns a NumPy array, e.g. Pillow, OpenCV, etc.
image = load_image()

# Simple left‑right flip – works on CPU, GPU or TPU.
flipped = pix.flip_left_right(image)

Because pix.flip_left_right is a pure JAX function you can also do:

import jax

# Compile once for maximum speed.
flipped = jax.jit(pix.flip_left_right)(image)

# Apply to a batch of images without writing a loop.
batch = image[None, ...]               # add a leading batch dimension
flipped_batch = jax.vmap(pix.flip_left_right)(batch)

Installation

  1. Install JAX first – follow the official JAX installation guide and pick the version that matches your CUDA/TPU setup.
  2. Install PIX from PyPI:
    pip install dm-pix
    (PIX itself is pure Python; it only needs the JAX runtime, which is why JAX isn’t listed as an automatic dependency.)

What can you do with PIX?

  • Basic geometric transforms: flips, rotations, crops, resizing.
  • Color‑space utilities (e.g., RGB↔HSV).
  • Convolution‑style filters that run efficiently on accelerators.
  • Any operation can be combined with JAX’s automatic differentiation, so you could, for example, optimise image‑processing pipelines as part of a larger learning system.

All examples live in the examples/ folder of the repository and are deliberately simple to illustrate how the functions compose with JAX’s jit, vmap, and pmap.

Testing & contribution

  • Run the test suite with pytest (install extra test dependencies via pip install -e "[test]").
  • Contributions are welcomed; see CONTRIBUTING.md for the workflow.

When to use PIX?

If you are already using JAX for model development and need fast, GPU‑friendly image preprocessing that can be compiled and parallelised alongside your model code, PIX provides a ready‑made, well‑tested set of utilities. It saves you from writing custom NumPy‑only code that would otherwise become a bottleneck when moving to accelerators.

TL;DR: PIX = JAX‑native image processing, fully jit‑able and parallelisable, install with pip install dm-pix after setting up JAX.