🧠 NeuroPack

NeuroPack is a research-oriented neural compression toolkit for 3D microscopy volumes. It adapts CompressAI-style learned image compression models to single-channel 3D TIFF stacks, encodes volumes into compact custom .bin bitstreams, and reconstructs decoded volumes as TIFF files.

---

✨ Key Features

• 3D learned image codec for single-channel volumetric TIFF data.
• Patch-based encoding for large volumes that do not fit into GPU memory as a single tensor.
• Custom binary stream format that stores model metadata, original volume size, bit depth, patch size, padding metadata, and entropy-coded patch latents.
• CUDA acceleration when PyTorch with CUDA is available.
• Batch encoding and decoding shell scripts that can be adapted to local data folders.
• Utilities inherited from CompressAI workflows for benchmarking, evaluation, plotting, and checkpoint entropy-model updates.

🗂️ Repository Layout

.
|-- codec.py                  # Main command-line encoder/decoder
|-- codec_func.py             # Importable Python encode/decode helpers
|-- encoding.sh               # Single-volume encoding example
|-- decoding.sh               # Single-volume decoding example
|-- multi_encoding.sh         # Batch encoding example
|-- multi_decoding.sh         # Batch decoding example
|-- requirements.txt          # Core runtime dependencies
|-- checkpoint/               # Example local checkpoint location, if present
|-- datasets/                 # Neuron TIFF dataset and patch sampling helpers
|-- layers/                   # 3D convolution, residual, GDN, and mask layers
|-- models/                   # 3D model definitions
|-- utils/                    # Benchmark, evaluation, plotting, model update tools
`-- zoo/                      # Model registry and architecture factory helpers

⚙️ Requirements

The main codec path is pinned for a CUDA 11.6-era PyTorch stack:

• Python 3.8 or 3.9 is recommended for the pinned dependencies.
• NVIDIA GPU and CUDA-capable PyTorch are recommended for practical inference.
• CPU execution is supported by the CLI, but large 3D volumes will be slow.
• 8-bit, single-channel TIFF stacks are the expected input format.

Core dependencies are listed in requirements.txt:

torch==1.13.1+cu116
torchvision==0.14.1+cu116
torchaudio==0.13.1
compressai==1.2.6
tifffile==2023.7.10
ipdb==0.13.13

Some utility modules require additional packages that are not pinned in requirements.txt, such as scipy, pillow, matplotlib, and pytorch-msssim.

---

🚀 Preparation

🧰 Installation

Create a fresh environment, then install the pinned runtime:

python -m venv .venv
source .venv/bin/activate
python -m pip install --upgrade pip
python -m pip install -r requirements.txt

If your CUDA or Python version differs from the pinned stack, install the matching PyTorch build first, then install the remaining packages manually:

python -m pip install compressai==1.2.6 tifffile==2023.7.10 ipdb==0.13.13

📥 Checkpoint and Sample Image

Encoding and decoding require a trained mbt2018-mean3D checkpoint:

checkpoint/checkpoint_best_loss.pth.tar

You can download the checkpoint and sample image data from the shared Google Drive folder. After downloading, place the checkpoint at the path above or update --load_model, and place the sample TIFF under a sample_name/ folder.

Pass the checkpoint path with --load_model for both encoding and decoding. The same checkpoint must be used to decode streams generated by the encoder.

📦 Expected Output Structure

The expected batch output structure is:

sample_name/
|-- bin_file/
|   `-- *.bin
`-- rec/
    `-- *.tif

---

⚡ Quick Start

🔐 Encode a 3D TIFF Volume

mkdir -p /path/to/sample_name/bin_file

python codec.py encode /path/to/sample_name/input_volume.tif \
  --model mbt2018-mean3D \
  --load_model checkpoint/checkpoint_best_loss.pth.tar \
  --quality 1 \
  --patch_size "64,128,128" \
  --cuda \
  -o /path/to/sample_name/bin_file/input_volume.bin

Target encoding result:

sample_name/
`-- bin_file/
    `-- input_volume.bin

Successful encoding prints the achieved bitrate and elapsed time:

0.105807 bpp | Encoded in 20.80s

🔓 Decode a Bitstream

mkdir -p /path/to/sample_name/rec

python codec.py decode /path/to/sample_name/bin_file/input_volume.bin \
  --load_model checkpoint/checkpoint_best_loss.pth.tar \
  --cuda \
  -o /path/to/sample_name/rec/input_volume.tif

Target decoding result:

sample_name/
`-- rec/
    `-- input_volume.tif

Successful decoding prints the model metadata embedded in the bitstream and the elapsed time:

Model: mbt2018-mean3D, metric: mse, quality: 1
Decoded in 31.65s

🖥️ Run Without CUDA

Omit --cuda to run on CPU:

mkdir -p /path/to/sample_name/bin_file

python codec.py encode /path/to/sample_name/input_volume.tif \
  --model mbt2018-mean3D \
  --load_model checkpoint/checkpoint_best_loss.pth.tar \
  --quality 1 \
  --patch_size "64,128,128" \
  -o /path/to/sample_name/bin_file/input_volume.bin

🔁 Batch Processing

The repository includes shell scripts for common local workflows:

• encoding.sh: encode one TIFF volume.
• decoding.sh: decode one .bin bitstream.
• multi_encoding.sh: traverse subfolders and encode matching *_R.tif files.
• multi_decoding.sh: traverse subfolders and decode .bin files into rec/.

These scripts contain absolute paths from the original development machines. Before running them, edit:

• INPUT_FILE or INPUT_FOLDER
• OUTPUT_FILE or OUTPUT_FOLDER
• MODEL_PATH
• PATCH_SIZE
• USE_CUDA

Then run:

bash encoding.sh
bash decoding.sh
bash multi_encoding.sh
bash multi_decoding.sh

---

🐍 Python API

For scripted use, codec_func.py exposes importable wrappers:

from codec_func import encode, decode

encode(
    input_path="/path/to/sample_name/input_volume.tif",
    model="mbt2018-mean3D",
    load_model="checkpoint/checkpoint_best_loss.pth.tar",
    metric="mse",
    quality=1,
    coder="ans",
    output="/path/to/sample_name/bin_file/input_volume.bin",
    cuda=True,
)

decode(
    input_path="/path/to/sample_name/bin_file/input_volume.bin",
    load_model="checkpoint/checkpoint_best_loss.pth.tar",
    coder="ans",
    output="/path/to/sample_name/rec/input_volume.tif",
    cuda=True,
)

Use the command-line interface when you need explicit patch-size control. The Python helper uses the fixed padding behavior implemented in codec_func.py and expects coder to be one of the entropy coders available in CompressAI.

🧾 Data Format Notes

• Input volumes are loaded with tifffile.imread.
• The main path expects a single-channel 3D array shaped like (depth, height, width).
• The current writer saves reconstructed volumes as 8-bit TIFF.
• The encoder normalizes uint8 TIFF data by 255; other integer bit depths are not currently handled as a first-class input format.
• The .bin format is project-specific and is intended to be decoded by this repository with a compatible checkpoint.

🧪 Utilities

The utils/ folder contains optional tools:

• utils/update_model: export checkpoints after updating entropy model CDFs.
• utils/eval_model: evaluate CompressAI image models on image datasets.
• utils/bench: benchmark traditional codecs such as JPEG, WebP, BPG, VTM, HM, AV1, and JPEG2000 when the required external binaries are installed.
• utils/plot: plot rate-distortion curves from JSON outputs.
• utils/medical3D: experimental 3D medical/neuron evaluation scripts.
• utils/video: video benchmark and evaluation helpers.

These utilities are not required for the basic 3D TIFF encode/decode workflow.

---

🛠️ Troubleshooting

❗ ModuleNotFoundError: No module named 'torch'

Install the PyTorch build that matches your Python and CUDA version. The pinned requirements target PyTorch 1.13.1+cu116.

💾 CUDA runs out of memory

Reduce --patch_size, for example:

--patch_size "32,64,64"

Smaller patches reduce memory pressure but may increase runtime and overhead.

🧯 Decoding opens an interactive debugger

The current development version of codec.py contains an ipdb.set_trace() call inside decode_image. Remove or comment that line for unattended batch decoding.

🔎 Batch scripts cannot find files

The shell scripts are templates with absolute paths. Replace every input, output, and checkpoint path with paths that exist on your machine.

⚠️ Current Limitations

• The primary maintained workflow is 3D TIFF volume compression with mbt2018-mean3D.
• The command-line code contains some legacy image/video branches from CompressAI that are not the main NeuroPack path.
• setup.py is currently empty, so use the project from the repository root rather than relying on package installation.
• A top-level LICENSE file is not present in this checkout. Several source files include upstream CompressAI/InterDigital BSD-style copyright notices.

---

📚 Citation and Acknowledgements

NeuroPack builds on the CompressAI ecosystem and adapts learned image compression components to 3D volumetric data. If you use this code in research, cite the relevant learned compression and CompressAI work, and add the project or model-specific citation used by your checkpoint.