google-deepmind/dm_pix

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

What is PIX?

PIX는 많은 현대적 머신러닝 연구를 뒷받침하는 고성능 NumPy 호환 프레임워크인 JAX를 기반으로 구축된 경량 이미지 처리 라이브러리입니다. PIX의 모든 함수(예: flip, resize, color conversion)는 jax.jit으로 컴파일하거나, jax.vmap으로 벡터화하거나, jax.pmap으로 여러 장치에서 실행할 수 있도록 작성되었습니다. 실제로 이는 이미지를 다른 JAX 배열과 동일하게 취급하여 신경망 코드에서 얻는 것과 동일한 GPU/TPU 가속을 얻을 수 있음을 의미합니다.

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)

pix.flip_left_right는 순수 JAX 함수이므로 다음과 같이 할 수도 있습니다:

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 자체는 순수 Python입니다. JAX 런타임만 필요하므로 JAX가 자동 의존성으로 나열되지 않습니다.)

What can you do with PIX?

  • 기본 기하학적 변환: flip, rotation, crop, resizing.
  • 색 공간 유틸리티(예: RGB↔HSV).
  • 가속기에서 효율적으로 실행되는 컨볼루션 스타일 필터.
  • 모든 작업은 JAX의 자동 미분 기능과 결합될 수 있으므로, 예를 들어 이미지 처리 파이프라인을 더 큰 학습 시스템의 일부로 최적화할 수 있습니다.

모든 예제는 저장소의 examples/ 폴더에 있으며, functions가 JAX의 jit, vmap, pmap과 어떻게 구성되는지 보여주기 위해 의도적으로 단순하게 작성되었습니다.

Testing & contribution

  • pytest를 사용하여 테스트 스위트를 실행합니다 (pip install -e "[test]"를 통해 추가 테스트 의존성을 설치하세요).
  • 기여는 언제나 환영합니다. 워크플로우는 CONTRIBUTING.md를 참조하세요.

When to use PIX?

이미 JAX를 사용하여 모델 개발을 하고 있으며, 모델 코드와 함께 컴파일 및 병렬화가 가능한 빠르고 GPU 친화적인 이미지 전처리가 필요하다면, PIX는 준비된, 잘 테스트된 유틸리티 세트를 제공합니다. 이는 가속기 사용 시 병목 현상이 될 수 있는 NumPy 전용 커스텀 코드를 작성해야 하는 번거로움을 덜어줍니다.

TL;DR: PIX = JAX 네이티브 이미지 처리, 완전히 jit‑able 및 병렬화 가능, JAX 설정 후 pip install dm-pix로 설치.