repository·master·Indexed 19 days ago
https://github.com/google-deepmind/dm_pixAn image processing library built on top of JAX, designed for high-performance image processing that can be optimized and parallelized using JAX transformations such as jit, vmap, and pmap. It provides a wide range of deterministic and stochastic image augmentations, color space conversions, depth and space transformations, N-dimensional linear interpolation, patch extraction, and image quality metrics like PSNR, SSIM, and MSE.
jax.jit, jax.vmap, and jax.pmap.The dm_pix library provides a wide range of image augmentation functions for deterministic and stochastic transformations.
Deterministic Augmentations:
adjust_brightness, adjust_contrast, adjust_gamma, adjust_hue, adjust_saturation, solarize.affine_transform, center_crop, elastic_deformation, flip_left_right, flip_up_down, pad_to_size, resize_with_crop_or_pad, rotate, rot90.gaussian_blur.Random (Stochastic) Augmentations:
random_brightness, random_contrast, random_crop, random_flip_left_right, random_flip_up_down, random_gamma, random_hue, random_saturation.To use PIX, import it as dm_pix. The library is designed to work seamlessly with JAX transformations like jax.jit, jax.vmap, and jax.pmap to provide optimized and parallelized image processing.
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)To run the provided examples, you must install the additional dependencies listed in the requirements file using pip.
$ pip install -r requirements_examples.txtYou can run the image augmentation examples locally by navigating to the examples/ directory and executing image_augmentation.py with Python.
$ cd examples/
$ python image_augmentation.pyPIX depends on JAX. Because JAX installation varies based on your CUDA version, you must install JAX first following the official JAX installation instructions with the appropriate accelerator support. Once JAX is installed, you can install PIX using pip.
$ pip install dm-pixAll PIX functions can take advantage of JAX's optimization and parallelization capabilities. You can use them as vanilla Python functions, or wrap them in jax.jit for compilation, jax.vmap for vectorization, or jax.pmap for multi-device parallelization.
Note: When using jax.vmap or jax.pmap, you typically need to add a leading dimension to your image (e.g., using image[np.newaxis, ...]) to represent the batch/device dimension.
import dm_pix as pix
import jax
import numpy as np
# Load an image into a NumPy array with your preferred library.
image = load_image()
# 1. Vanilla Python function.
flip_left_right_image = pix.flip_left_right(image)
# 2. `jax.jit`ed function (compilation).
flip_left_right_image = jax.jit(pix.flip_left_right)(image)
# Add a leading dimension for parallelization (vmap/pmap)
image = image[np.newaxis, ...]
# 3. `jax.vmap`ed function (vectorization).
flip_left_right_image = jax.vmap(pix.flip_left_right)(image)
# 4. `jax.pmap`ed function (multi-device parallelization).
flip_left_right_image = jax.pmap(pix.flip_left_right)(image)To test your development environment, you can install the test dependencies and run pytest.
To run in an isolated virtual environment, use the provided ./test.sh script.
# Install test dependencies and run pytest
$ pip install -e ".[test]"
$ python -m pytest [-n <NUMCPUS>] dm_pix
# Or use the utility script for an isolated environment
$ ./test.shYou can install dm-pix using pip.
!pip install dm-pixextract_patches to divide an image into smaller sub-regions or patches.Use the following functions to transform images between different color representations:
rgb_to_hsl, rgb_to_hsv, rgb_to_grayscale.hsl_to_rgb, hsv_to_rgb.Transform the spatial arrangement of pixels using:
depth_to_spacespace_to_depth