Metadata-Version: 2.4
Name: medtokenizers
Version: 0.1.0
Summary: Tokenization modules for volumetric medical imagery.
Author: Liam Chalcroft
License-Expression: MIT
Project-URL: Homepage, https://github.com/liamchalcroft/medtokenizers
Project-URL: Documentation, https://liamchalcroft.github.io/medtokenizers/
Project-URL: Repository, https://github.com/liamchalcroft/medtokenizers
Project-URL: Issues, https://github.com/liamchalcroft/medtokenizers/issues
Project-URL: Changelog, https://github.com/liamchalcroft/medtokenizers/blob/main/CHANGELOG.md
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Science/Research
Classifier: Topic :: Scientific/Engineering :: Medical Science Apps.
Classifier: Topic :: Scientific/Engineering :: Image Recognition
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Operating System :: POSIX :: Linux
Classifier: Operating System :: MacOS :: MacOS X
Requires-Python: >=3.10
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: torch>=2.0
Requires-Dist: numpy>=1.23
Requires-Dist: einops>=0.6
Requires-Dist: jaxtyping>=0.2.20
Requires-Dist: beartype>=0.15.0
Requires-Dist: tqdm>=4.66.0
Requires-Dist: torchvision>=0.23.0
Requires-Dist: huggingface-hub>=0.19.0
Requires-Dist: transformers>=4.30.0
Requires-Dist: medrs>=0.2.0
Requires-Dist: lpips>=0.1.4
Requires-Dist: nibabel>=5.3.3
Provides-Extra: examples
Requires-Dist: accelerate>=0.25.0; extra == "examples"
Requires-Dist: wandb>=0.16.0; extra == "examples"
Requires-Dist: medmnist>=3.0.0; extra == "examples"
Requires-Dist: matplotlib>=3.10.8; extra == "examples"
Requires-Dist: mgzip>=0.2.1; extra == "examples"
Requires-Dist: scipy>=1.10; extra == "examples"
Requires-Dist: scikit-image>=0.22; extra == "examples"
Requires-Dist: safetensors>=0.4.0; extra == "examples"
Provides-Extra: test
Requires-Dist: medtokenizers[examples]; extra == "test"
Requires-Dist: pytest; extra == "test"
Requires-Dist: pytest-cov; extra == "test"
Provides-Extra: dev
Requires-Dist: medtokenizers[test]; extra == "dev"
Requires-Dist: ruff>=0.4.0; extra == "dev"
Provides-Extra: docs
Requires-Dist: sphinx>=7.0; extra == "docs"
Requires-Dist: sphinx-rtd-theme>=1.3; extra == "docs"
Requires-Dist: sphinx-autodoc-typehints>=2.0; extra == "docs"
Requires-Dist: sphinxcontrib-mermaid>=0.9; extra == "docs"
Requires-Dist: myst-parser>=2.0; extra == "docs"
Provides-Extra: cloud
Requires-Dist: medtokenizers[training]; extra == "cloud"
Requires-Dist: python-dotenv>=1.0.0; extra == "cloud"
Provides-Extra: training
Requires-Dist: accelerate>=0.25.0; extra == "training"
Requires-Dist: wandb>=0.16.0; extra == "training"
Provides-Extra: all
Requires-Dist: medtokenizers[cloud,dev,docs,training]; extra == "all"
Dynamic: license-file

# medtokenizers

`medtokenizers` compresses 2D and 3D medical images into latents or discrete
token grids. It provides one encoder-decoder backbone with interchangeable
quantization heads (VQ, LFQ, FSQ, Residual FSQ, VAE, AE), so the quantizer can
be treated as a controlled variable rather than a fixed preprocessing choice.
It also provides shared reconstruction metrics and dataset tokenization
utilities.

This library accompanies the paper *Tokenizer-Generator Coupling in Medical
Image Generation*, which uses it for the tokenizer half of a factorial study.
See [Citation](#citation).

## Install

```bash
pip install -e ".[test]"
```

Requires Python 3.10+ and PyTorch 2.0+. Tested on Linux and macOS. A CUDA GPU
is needed for 3D training, but everything below runs on CPU.

For a CPU-only install, take PyTorch from its CPU index first so `torch` and
`torchvision` come from the same build:

```bash
pip install torch torchvision --index-url https://download.pytorch.org/whl/cpu
pip install -e ".[test]"
```

## Run it

Swap the quantizer, keep everything else fixed:

```python
import torch

from medtokenizers import DiscreteTokenizer

x = torch.randn(1, 1, 64, 64)

for quantizer in ["VQ", "FSQ", "RESFSQ"]:
    model = DiscreteTokenizer(
        dim=2, quantizer=quantizer, resolution=64, spatial_compression=8
    ).eval()
    with torch.no_grad():
        tokens = model.tokenize(x)  # (1, 8, 8) integer codes
        recon = model.detokenize(tokens)  # (1, 1, 64, 64)
```

LFQ is the one head that needs two further arguments, because its codebook is
defined by its binary dimension rather than inferred:

```python
model = DiscreteTokenizer(
    dim=2,
    quantizer="LFQ",
    resolution=64,
    spatial_compression=8,
    codebook_size=1024,
    codebook_dim=10,
)
```

The continuous heads live on a separate class, selected by `formulation`
instead of `quantizer`, and return a real-valued latent rather than indices:

```python
from medtokenizers import ContinuousTokenizer

model = ContinuousTokenizer(
    dim=2, formulation="VAE", resolution=64, spatial_compression=8
).eval()
with torch.no_grad():
    latent = model.tokenize(x)  # (1, 4, 8, 8)
    recon = model.detokenize(latent)
```

At evaluation time `ContinuousTokenizer.encode` returns the posterior mean, so
latents extracted with `model.eval()` are deterministic.

Set `dim=3` for volumes; the same calls take `(B, C, H, W, D)` tensors.

## Volumetric inference

A simulated T1-weighted brain volume ships inside the package (from the BrainWeb
Simulated Brain Database, with provenance in
[src/medtokenizers/assets/README.md](src/medtokenizers/assets/README.md)), so the
example runs without a download and works from an installed wheel.
`example_volume_path()` resolves it. Random initialisation is enough for a smoke
test; load your own weights for real reconstructions.

```bash
python examples/inference_on_brain.py --model-type maisi
```

Large volumes are reconstructed with Gaussian-weighted sliding windows:

```python
import medrs
import torch

from medtokenizers import MAISITokenizer, example_volume_path

img = medrs.load(str(example_volume_path()))
volume = img.to_torch_with_dtype_and_device(dtype=torch.float32)
volume = volume.unsqueeze(0).unsqueeze(0)  # (H, W, D) -> (1, 1, H, W, D)
volume = (volume - volume.min()) / (volume.max() - volume.min() + 1e-8)

model = MAISITokenizer().eval()
with torch.no_grad():
    recon = model.reconstruct(volume, roi_size=(96, 96, 64), overlap=0.5)
```

Sliding-window reconstruction over a full 1 mm volume is memory-hungry. Reduce
`roi_size` or `overlap` if you are not on a GPU.

## Tokenizers

| Class | Selected by | Output |
| --- | --- | --- |
| `DiscreteTokenizer` | `quantizer="VQ" \| "FSQ" \| "LFQ" \| "RESFSQ"` | integer token grid |
| `ContinuousTokenizer` | `formulation="VAE" \| "AE"` | real-valued latent |
| `MAISITokenizer` | fixed MAISI configuration | real-valued latent |
| `TiTokTokenizer` | 1D transformer tokenizer | fixed-length sequence of K tokens |
| `RAETokenizer` *(experimental)* | frozen foundation-model encoder | real-valued latent |

Adding a quantizer means subclassing `BaseQuantizer` (two abstract methods,
`forward` and `indices_to_codes`) and wiring it into `DiscreteTokenizer`. The
encoder-decoder and the tokenization I/O are unchanged by the choice.

### TiTok

TiTok encodes an image or volume into a fixed-length 1D sequence of K tokens.
The encoder concatenates patch tokens with K learnable latent tokens and keeps
only the latter; the decoder reconstructs from quantized latents plus mask
tokens. Input shape must equal `resolution`, and each spatial dimension must be
divisible by `patch_size`.

```python
from medtokenizers import TiTokTokenizer

tokenizer = TiTokTokenizer(
    dim=2,
    in_channels=1,
    out_channels=1,
    resolution=64,
    patch_size=16,
    num_tokens=32,
    num_embeddings=1024,
    hidden_dim=256,
)
indices, codes, quant_loss = tokenizer.encode(torch.randn(2, 1, 64, 64))
recon = tokenizer.decode(codes)
```

TiTok's original two-stage proxy-code warmup and decoder fine-tuning are not
wired into the training scripts; implement that externally if you want the
published recipe.

## Evaluation

```python
from medtokenizers import compute_lpips, compute_psnr, compute_ssim

psnr = compute_psnr(reference, reconstruction)
ssim = compute_ssim(reference, reconstruction)
lpips = compute_lpips(reference, reconstruction)
```

`TokenizerEvaluator` runs these over a loader, together with perplexity and
codebook utilisation for discrete models:

```python
from medtokenizers import TokenizerEvaluator, load_tokenizer

model = load_tokenizer("./path/to/checkpoint")
results = TokenizerEvaluator(model, device="cuda").evaluate(test_loader)
```

No pretrained checkpoints ship with this repository.

## Ecosystem

`medtokenizers` is one of three packages:

- [medrs](https://github.com/liamchalcroft/med-rs): medical-image I/O
  (NIfTI/DICOM) and preprocessing primitives.
- [medtokenizers](https://github.com/liamchalcroft/medtokenizers): tokenizers
  that compress scans into latents or token grids *(this package)*.
- [medlatents](https://github.com/liamchalcroft/medlatents): generative models
  (autoregressive, MaskGIT, diffusion, flow matching, Bayesian flow) that learn
  over those tokens.

```text
images -> medrs (load) -> medtokenizers (tokenize) -> medlatents (generate)
       -> medtokenizers (detokenize) -> images
```

`scripts/tokenize_dataset.py` writes a directory of per-split `.npz` files plus
`metadata.json`, with key `codes` (int16) for discrete models and `latents`
(float16) for continuous ones.

## Contributing

See [CONTRIBUTING.md](CONTRIBUTING.md) for development setup and the
pull-request workflow, and [CODE_OF_CONDUCT.md](CODE_OF_CONDUCT.md) for
community standards.

## Citation

```bibtex
@article{chalcroft2026coupling,
  author = {Chalcroft, Liam},
  title  = {Tokenizer--Generator Coupling in Medical Image Generation},
  year   = {2026},
  note   = {arXiv preprint, to appear}
}

@software{chalcroft_medtokenizers,
  author  = {Chalcroft, Liam},
  title   = {{medtokenizers}: Continuous and discrete tokenizers for volumetric medical imaging},
  url     = {https://github.com/liamchalcroft/medtokenizers},
  version = {0.1.0}
}
```

Machine-readable metadata is in [CITATION.cff](CITATION.cff).

## License

MIT, see [LICENSE](LICENSE).

Some modules derive from third-party projects: CompVis (Stable Diffusion,
taming-transformers) and lucidrains `vector-quantize-pytorch` under MIT, and
NVIDIA Cosmos-Tokenizer and MONAI/MAISI under Apache-2.0. Their terms are
preserved in [THIRD_PARTY_NOTICES.md](THIRD_PARTY_NOTICES.md) and
[NOTICE](NOTICE). Two things are outside the MIT licence: NVIDIA MAISI
pretrained *weights*, if you use them, remain under NSCLv1; and the bundled
bundled BrainWeb volume is data carrying its own citation requirement.
