INIF

This repo is the offical codebase for the manuscript 'Implicit Neural Image Field for Biological Microscopy Image Compression'.

Introduction

Implicit neural image field (INIF) is a SOTA paradigm for microscopy image compression, characterized by its effectiveness and flexibility. It adopts a learned optimizer and leverages common codecs for faster compression, while integrating application-specific guidance for improved compression quality and trustworthiness.

System Requirements

Operating System

Note

Currently only Linux System is supported. The test enviornment is Ubuntu 20.04, RTX3090 GPU, and CUDA 11.8. For local installation, please make sure the CUDA version is compatible with the jax and tensorflow version. Otherwise we strongly suggest to tryout the tutorial in

Dependencies

Framework

• jax==0.4.26[cuda12_pip]
• tensorflow==2.8.4 (Nvidia-GPU version)

pip

• dm-haiku==0.0.13.dev0 (Nvidia-GPU version)
• optax==0.2.2
• ipywidgets==8.1.2
• tifffile==2024.2.12
• opencv-python==4.9.0.80
• matplotlib==3.8.4
• pandas==2.2.1
• learned-optimization==0.0.1 (lopt)

Installation

Note

The typical install time is under 10 minutes. (Test it on Google Colab's free GPU)

From source:

1. clone this repository:

   git clone https://github.com/daviddaiiiii/INIF.git

2. Create the conda environment and activate:

   conda deactivate
   conda create -n INIF python=3.10
   conda activate INIF

3. Run the setup.sh file to install automatically:

   chmod +x setup.sh
   ./setup.sh

   These codes would:
   1. Give the script execution permissions
   2. Install the dependencies via pip

Tip

In case of issues when downloading the weight of learned optimizers automatically from Google Cloud, below are some workaround:

Some workaround:

1. We provide the default weight and continue trained weight in '../INIF/Learned_optimizer_weight'

2. go to '../anaconda3/envs/INIF/lib/pythonX.XX/site-packages/learned_optimization/research/general_lopt/pretrained_optimizers.py'

   💡 ctrl+click on 'from learned_optimization.research.general_lopt import pretrained_optimizers' to jump to the file

3. replace the path in line143 to '../INIF/Learned_optimizer_weight

   _pretrain_no_config_root = 'gs://gresearch/learned_optimization/pretrained_lopts/no_config/'

Docker:

1. Install docker if you don't have it: get-docker

Important

To utilize GPU, nvidia-docker package is also required.

2. retrieve the docker image from the dockerhub:

   docker pull pkuhmi/inif_env:latest

3. start a container:

   docker run --gpus all -it --rm --shm-size 16G --ulimit memlock=-1 -v ./INIF:/INIF/ --name inif pkuhmi/inif_env:latest

   where:

   • --gpus: use the gpu
   • -it: interact with the container
   • --rm: remove the container after exit
   • --shm-size: set the shared memory size to avoid memory issues
   • --ulimit memlock=-1: remove mem lock limit
   • -v: mount the current directory to the container
   • --name: name the container

Tip

We recommend to use the VSCode Docker extension to attach the container to the VSCode workspace and run the tutorial.

Quick Start 🚀

We provide a google colab notebook that demonstrates how to use INIF for the general image compression:

You can also try it out locally by playing around with the compression & decompression.

Note

The typical runtime for traning depends on the raw data and device used. In the tutorial, it would take approximately 10 minutes to train the model on a Google Colab's free GPU.