Package: dm-pix
PIXPIX is an image processing library inJAX, forJAX.OverviewJAXis a library resulting from the union ofAutogradandXLAfor
high-performance machine learning research. It providesNumPy,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
throughjax.jit,jax.vmapandjax.pmap.InstallationPIX 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 inrequirements.txt, although it is technically
listed for reference, but commented.First, followJAX installation instructionsto install JAX with the relevant
accelerator support.Then, install PIX usingpip:$pipinstalldm-pixQuickstartTo usePIX, you just need toimport dm_pix as pixand use it right away!For example, let's assume to have loaded the JAX logo (available inexamples/assets/jax_logo.jpg) in a variable calledimageand we want to flip
it left to right.All it's needed is the following code!importdm_pixaspix# 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!All the functions in PIX can bejax.jited,jax.vmaped andjax.pmaped, so all the following functions can take advantage of
optimization and parallelization.importdm_pixaspiximportjax# 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 ofpix.flip_left_rightis the same (up to the accelerator floating point
accuracy)!ExamplesWe have a few examples in theexamples/folder. They are not much
more involved then the previous example, but they may be a good starting point
for you!TestingWe 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_testsuffix, and can be executed usingpytest.If you already have PIX installed, you just need to install some extra
dependencies and runpytestas follows:$pipinstall-rrequirements_tests.txt
$python-mpytest[-n<NUMCPUS>]dm_pixIf you want an isolated virtual environment, you just need to run our utilitybashscript as follows:$./test.shCiting PIXThis repository is part of theDeepMind JAX Ecosystem, to cite PIX please use
theDeepMind JAX Ecosystem citation.Contribute!We are very happy to accept contributions!Please read ourcontributing guidelinesand send us PRs!