CASTOR

Welcome to the code repository for projects related to the CArdiac SegmenTation with cOnstRaints (CASTOR) project.

Publications

Description

This is a project that constrains the predictions of automatic cardiac segmentation a posteriori to guarantee useful properties, i.e. anatomical validity and temporal consistency.

To help you follow along with the organization of the repository, here is a summary of each major package's purpose:

• apps: interactive applications, either graphical or command line, that help to inspect data and/or results.

• results: API and executable scripts for processing results during the evaluation phase.

• vital: a separate repository (included as a git submodule), of generic PyTorch modules, losses and metrics functions, and other tooling (e.g. image processing, parameter groups) that are commonly used. Also contains the code for managing specialized medical imaging datasets, e.g. ACDC, CAMUS.

How to Run

Install

First, download the project's code:

# clone project
git clone --recurse-submodules https://github.com/vitalab/castor.git

Next you have to install the project and its dependencies. The project's dependency management and packaging is handled by poetry so the recommended way to install the project is in a virtual environment (managed by your favorite tool, e.g. conda, virtualenv, poetry, etc.), where poetry is installed. That way, you can simply run the command:

poetry install

from the project's root directory to install it in editable mode, along with its regular and development dependencies. This command also takes care of installing the local vital submodule dependency in editable mode, so that you can edit the library and your modifications will be automatically taken into account in your virtual environment.

Note When a poetry.lock file is available in the repository, poetry install will automatically use it to determine the versions of the packages to install, instead of resolving anew the dependencies in pyproject.toml. When no poetry.lock file is available, the dependencies are resolved from those listed in pyproject.toml, and a poetry.lock is generated automatically as a result.

Warning Out-of-the-box, poetry offers flexibility on how to install projects. Packages are natively pip-installable just as with a traditional setup.py by simply running pip install <package>. However, we recommend using poetry because of an issue with pip-installing projects with relative path dependencies (the vital submodule is specified using a relative path). When the linked issue gets fixed, the setup instructions will be updated to mention the possibility of using pip install ., if one wishes to avoid using poetry entirely.

To test that the project was installed successfully, you can try the following command from the Python REPL:

# now you can do:
from castor import Whatever

Note The instructions above for setting up an environment are for general purpose/local environments. For more specific use cases, e.g. on DRAC clusters, please refer to the installation README.

Warning All following commands in this README (and other READMEs for specific packages), will assume you're working from inside the virtual environment where the project is installed.

Use pretrained models for inference

Comet model registry

The pretrained baseline models listed in the following sections are provided as public models from our Comet model registry. Our code can handle how to download and load these models transparently, provided you configure a Comet API key.

In case you don't want to use our integration with Comet's API to automatically download the models in the background (or if you want to run the code in environments where there is no internet access, e.g. on HPC clusters), there is the option to manually download the checkpoints. For instructions on how do this, we refer you to Comet's documentation on the subject.

Note In case you manually download the pretrained models, you should replace the name of the model in the examples below with the path of the .ckpt file (the .ckpt file can be extracted from the archive you get when downloading a model from the Comet registry).

Automatic segmentation using baseline models

Available baseline segmentation models (for more information on how to use the Comet model registry, refer to this section of the README):

• ENet: "nathanpainchaud/echo-enet"
• U-Net: "nathanpainchaud/echo-camusunet"

In the following cell, you will find a minimal working example of how to load a pretrained model and use it to predict the segmentation on a new batch of data.

from vital.utils.saving import load_from_checkpoint

# Load a pretrained model from the publicly available models in nathanpainchaud's Comet model registry
pretrained_model_name = "nathanpainchaud/echo-enet" # Can be replaced by any of the pretrained segmentation models listed above
model = load_from_checkpoint(checkpoint=pretrained_model_name)

# The result of the `load_from_checkpoint` call is an instance of the `vital.tasks.segmentation.SegmentationTask` class
# which we can simply use as callable to predict segmentations on a batch of images
your_image_batch: torch.Tensor # Tensor of shape (N, H, W) where N is the batch dimension
predicted_segmentation = model(your_image_batch)

Post-processing segmentations for temporal consistency

Available segmentation autoencoders (for more information on how to use the Comet model registry, refer to this section of the README):

• Cardiac AR-VAE: To be published at a later date

In the following cell, you will find a minimal working example of how to load a pretrained model and use it to post-process a new sequence of 2D segmentations for temporal consistency.

from vital.data.camus.utils.process import TEDTemporalRegularization

# Instantiate the class that handles temporal consistency post-processing over segmentations, using a pretrained
# cardiac AR-VAE from the publicly available models in nathanpainchaud's Comet model registry as a backbone
# NOTE: If you use the provided pretrained AR-VAE, your segmentations should label the left ventricle as 1 and the
#       myocardium as 2.
pretrained_model_name = <MODEL_NAME> # Can be replaced by any of the pretrained autoencoder models listed above
temporal_regularization = TEDTemporalRegularization(autoencoder=pretrained_model_name)

# The `TEDTemporalRegularization` is an object that can be used as a callable to perform temporal consistency
# post-processing. Since the result of the post-processing returns both the segmentation and the encoding in the latent
# space, we have to access the "post_mask" key to get the segmentation itself
your_image_sequence: torch.Tensor # Tensor of shape (N, H, W) where N is the temporal dimension
postprocessed_segmentation = temporal_regularization(your_image_sequence)["post_mask"]

Training the models yourself

Data

Navigate to the data folder for either the ACDC or CAMUS dataset and follow the instructions on how to setup the datasets:

• ACDC instructions
• CAMUS instructions

Configuring a Run

This project uses Hydra to handle the configuration of the castor runner script. To understand how to use Hydra's CLI, refer to its documentation. For this particular project, preset configurations for various parts of the castor runner pipeline are available in the config package. These files are meant to be composed together by Hydra to produce a complete configuration for a run.

Below we provide examples of how to run some basic commands using the Hydra CLI:

# list generic trainer options and datasets on which you can train
castor-runner -h

# select high-level options of task to run, and architecture and dataset to use
castor-runner task=<TASK> task/model=<MODEL> data=<DATASET>

# display the available configuration options for a specific combination of task/model/data (e.g Enet on CAMUS)
castor-runner task=segmentation task/model=enet data=camus -h

# train and test a specific system (e.g beta-VAE on CAMUS)
castor-runner task=autoencoder task/model=beta-vae data=camus data.dataset_path=<DATASET_PATH> [optional args]

# test a previously saved system (e.g. beta-VAE on CAMUS)
castor-runner task=autoencoder task/model=beta-vae data=camus data.dataset_path=<DATASET_PATH> \
  ckpt=<CHECKPOINT_PATH> train=False

# run one of the fully pre-configured 'experiment' from the `config/experiment` folder (e.g. Enet on CAMUS)
castor-runner +experiment=camus/enet

To create your own pre-configured experiments, like the one used in the last example, we refer you to Hydra's own documentation on configuring experiments.

Tracking experiments

By default, Lightning logs runs locally in a format interpretable by Tensorboard.

Another option is to use Comet to log experiments, either online or offline. To enable the tracking of experiments using Comet, simply use one of the pre-built Hydra configuration for Comet. The default configuration is for Comet in online mode, but you can use it in offline mode by selecting the corresponding config file when launching the castor runner script:

castor-runner logger=comet/offline ...

To configure the Comet API and experiment's metadata, Comet relies on either i) environment variables (which you can set in a .env that will automatically be loaded using python-dotenv) or ii) a .comet.config file. For more information on how to configure Comet using environment variables or the config file, refer to Comet's configuration variables documentation.

An example of a .comet.config file, with the appropriate fields to track experiments online, can be found here. You can simply copy the file to the directory of your choice within your project (be sure not to commit your Comet API key!!!) and fill the values with your own Comet credentials and workspace setup.

Note No change to the code is necessary to change how the CometLogger handles the configuration from the .comet.config file. The code simply reads the content of the [comet] section of the file and uses it to create a CometLogger instance. That way, you simply have to ensure that the fields present in your configuration match the behavior you want from the CometLogger integration in Lighting, and you're good to go!

How to Contribute

Environment Setup

When installing the dependencies using poetry install as described above, the resulting environment is already fully configured to start contributing to the project. There's nothing to change to get coding!

Version Control Hooks

Before first trying to commit to the project, it is important to setup the version control hooks, so that commits respect the coding standards in place for the project. The .pre-commit-config.yaml file defines the pre-commit hooks that should be installed in any project contributing to the vital repository. To setup the version control hooks, run the following command:

pre-commit install

Note In case you want to copy the pre-commit hooks configuration to your own project, you're welcome to :) The configuration for each hook is located in the following files:

• isort: pyproject.toml, [tool.isort] section
• black: pyproject.toml, [tool.black] section
• flake8: setup.cfg, [flake8] section

However, be advised that isort must be configured slightly differently in each project. The src_paths tag should thus reflect the package directory name of the current project, in place of vital.

References

If you find this code useful, please consider citing the paper implemented in this repository relevant to you from the list below:

@article{painchaud_echocardiography_2022,
    title = {Echocardiography {Segmentation} {with} {Enforced} {Temporal} {Consistency}},
    doi = {10.1109/TMI.2022.3173669},
	journal = {IEEE Transactions on Medical Imaging},
	author = {Painchaud, N. and Duchateau, N. and Bernard, O. and Jodoin, P.-M.},
	year = {2022},
}

@article{painchaud_cardiac_2020,
	title = {Cardiac {Segmentation} {With} {Strong} {Anatomical} {Guarantees}},
	volume = {39},
	copyright = {All rights reserved},
	issn = {1558-254X},
	doi = {10.1109/TMI.2020.3003240},
	number = {11},
	journal = {IEEE Transactions on Medical Imaging},
	author = {Painchaud, N. and Skandarani, Y. and Judge, T. and Bernard, O. and Lalande, A. and Jodoin, P.-M.},
	month = nov,
	year = {2020},
	pages = {3703--3713},
}