University of Cambridge Computational Image Analysis (CIA) groups AI tools for image reconstruction (e.g. tomographic reconstruction), LION (Learned Iterative Optimization Networks)
The library is a place to gather and provide resources for image reconstruction, such as Computed Tomography reconstruction, using machine learning tools implemented in PyTorch.
It contains data dowloading and pre-processing, torch DataSet classes for the existing data, and Experiments class with implementation of default realistic experiments, several data-driven methods and models, and features to enhance reproduciblity.
WARNING Building in progress. This is a tool in development in very early stages. Many things are bound to fail and many more are bound to change. If you want to help with development, send an email to Ander Biguri and/or open and issue or a discussion.
Install:
git clone https://github.com/CambridgeCIA/LION.git
cd LION
git submodule update --init --recursive # Legacy: for MSD_pytorch_
conda env create --file=env_base.yml --name=lion # You can change 'lion' to a different env name
conda activate lion
pip install .With mamba, use mamba env create --file=env_base.yml --name=lion and mamba activate lion instead; the rest is the same.
Everything compiled, torch included, comes from conda (env_base.yml), so all of it uses the same C++ runtime. Don't install torch with pip on top: a pip torch loads the system's C++ runtime, and conda's libraries (scipy, ...) then fail with CXXABI_... not found on machines with an older system. Create the environment fresh from env_base.yml rather than updating an old one.
- CUDA version: set by
cuda-versioninenv_base.yml(12.9). Any CUDA 12 driver runs it (nvidia-smishows the driver's CUDA version). - No GPU on the machine where you create the environment (e.g. a login node): conda cannot see a CUDA driver and refuses
pytorch-gpu. Tell it which CUDA to assume:CONDA_OVERRIDE_CUDA=12.9 conda env create --file=env_base.yml --name=lion.
If you would like to contribute to the development of LION, you can replace the last line of the above set of commands by
pip install -e ".[dev]"
to make the installation editable (i.e. changes you make to the source will be visible when you restart the REPL or start a new Python process) and include additional development dependencies like pre-commit and pytest-xdist.
Afterwards, install the pre-commit hooks for auto-formating your commits. We use Black for code formatting (see .pre-commit-config.yaml):
pre-commit install --hook-type pre-commit --hook-type post-merge
Set the device once, at the top of your script, and LION uses it wherever no device is given explicitly (solvers, loading models, datasets, noise simulation):
import torch
torch.cuda.set_device(2) # or: torch.set_default_device("cuda:2")CT operators (tomosipo/astra) run on the device of the tensors they are given. On machines with several GPUs where only one may be used, the most robust option is to make only that GPU visible to the process, so nothing (LION, PyTorch or astra) can use another:
CUDA_VISIBLE_DEVICES=2 python my_script.pySeveral DataSets are currently supported by LION. To point the package to the correct storage location use export LION_DATA_PATH=/path/to/Datasets. For convenience, this can be added to your .bashrc or similar file.
NOTE: If you are part of CIA at DAMTP, you already have access to these datasets and you don't need to follow any instruction to dowload them. They should be accessible using export LION_DATA_PATH=/store/LION/datasets
The supported Datasets are 2DeteCT, LIDC-IDRI,
Read more about the Datasets and how to download them here
LION supports all types of data-driven methods for CT reconstructions. They can, as a general taxonomy, be described as:
- Post-Processing methods: a "denoising" network. Takes a noisy recon and cleans it.
- Iterative Unrolled methods: Uses the operator to imitate iterative recon algorithms, but has learned parts.
- Learned regularizer: Explicitly learned regularization functions.
- Plug-and-Play (PnP): Implicit learned regularization, a regularization optimization step is learned, rather than an explicit one.
Folders for each of these exist in LION/models. An extra folder for standard CNNs is also available.
Read more about which models are available in each class here
Read more about which papers are implemented in LION here
Originally designed for tomographic reconstruction tasks, specifically Computed Tomography (CT), LION has been extended to work with more operators. Read more about operators here
Read developers.md
Run the tests with:
pytest pytests/Tests that need CUDA (cuda marker) or tomosipo/astra (tomosipo marker) are skipped automatically on machines without them (e.g. Macs). The folder is called pytests rather than tests because tomosipo installs a top-level tests package. The same tests run on every pull request (.github/workflows/tests.yml).
LION is released under the GNU General Public License v3 (GPL-3), see LICENSE.txt. It has to be: its CT backend, ASTRA and tomosipo, are GPL-3.