SHAPR

For a short video introducing our topological loss function and SHAPR please see:

@inproceedings{Waibel22a,
  author        = {Dominik J. E. Waibel and Scott Atwell and Matthias Meier and Carsten Marr and Bastian Rieck},
  title         = {Capturing Shape Information with Multi-Scale Topological Loss Terms for 3D Reconstruction},
  year          = {2022},
  booktitle     = {Medical Image Computing and Computer Assisted Intervention~(MICCAI)},
  archiveprefix = {arXiv},
  eprint        = {2203.01703},
  primaryclass  = {cs.CV},
  repository    = {https://github.com/marrlab/SHAPR_torch},
  pubstate      = {inpress},
}

For more information, please refer to our preprint on bioRxiv here.

SHAPR consists of a 2D encoder, which embeds 2D images into a latent space, and a 3D decoder, which reconstructs 3D shapes from the latent space representation. To train SHAPR we segment 3D microscopy images (we show an exemplary single red blood cell). We pair a 2D segmentation with the microscopy image of the same slice to enter the encoder as input. During supervised training (Fig. 1, step 1), we minimize the reconstruction loss (see Methods), which is the sum of the Dice loss and the binary cross entropy loss between the 3D segmentations and SHAPR predictions. For an input image of 64 x 64 pixels, we provide the pixel sizes for each layer in the gray boxes and the filter sizes on top of each box. In the second step, we fine-tune SHAPR by adding a discriminator. The discriminator is trained to differentiate between SHAPR output and ground truth segmentation and minimize the adversarial loss. It thereby challenges SHAPR to output realistic 3D objects.

Given a predicted object and a 3D ground truth object, we calculate topological features using (cubical) persistent homology, obtaining a set of persistence diagrams. Each point in a persistence diagram denotes the birth and death of a topological feature of some dimension of the given object. We compare these diagrams using $\mathcal{L}_T$, our topology-based loss, and weight this with existing loss terms like binary cross entropy (BCE) and Dice.

Installation

Fork the repository.

$ cd SHAPR_torch
$ pip3 install -e .

Installation using a virtual environment

We would recommend to first set a virtual environment and then install the package:

$ cd SHAPR
$ python3 -m venv .venv_shapr
$ source .venv_shapr/bin/activate
(.venv_shapr) $ pip3 install -e .

By activating the virtual environment your shell’s prompt will be changed in order to show what virtual environment you’re using.

You can deactivate a virtual environment by:

(.venv_shapr) $ deactivate

We can also use the virtual environment as a kernel for Jupyter Notebook. First you should install ipykernel package when the virtual environment is activated:

(.venv_shapr) $ pip3 install ipykernel

We need to manually add our virtual environment as a kernel to Jupyter Notebook:

(.venv_shapr) $ python -m ipykernel install --name=.venv_shapr

Now by opening the Jupyter-Notebook you have the option to select the .venv_shapr as the kernel.

Running SHAPR

Please find an example of how to run SHAPR from a jupyter notebook in SHAPR_torch/docs/jupyter notebook/Run SHAPR from notebook.ipynb.

You can also run SHAPR using a params.json file, which is provided in SHAPR_torch/docs/sample/params.json.

Setting parameters

To run SHAPR you should set the following parameters: Setting parameters are:

• path: path to a folder that includes three subfolder of
  1. obj: containing the 3D groundtruth segmentations,
  2. mask: containing the 2D masks,
  3. image: containing the images from which the 2D masks were segmented (e.g. brightfield).
• result_path: path to a folder for saving the results of predictions.
• pretrained_weights_path: path to a folder for saving and reloading pretrain model
• random_seed: seed for random generator in order to keep the results reproducible.

The setting parameters are read from the settings object. You may change the setting parameters by directly changing their default values in a SHAPR_torch/params.json file or simply package API like:

> from shapr import settings
> settings.path = "a/new/path"

We have added an example of a params.json file to SHAPR/docs/sample/params.json. If you want to use it, please adapt the paths to your project and copy the params.json to SHAPR_torch/params.json, then execute SHAPR_torch/shapr/run_train_script.py. You can also print all the parameters and their values using print:

> print(settings)
------ settings parameters ------
path: "path value"
result_path: "result_path value"
pretrained_weights_path: "pretrained_weights_path value"
random_seed: 0

Running specific parts of the training loop individually

You can run the training and evaluation on the test sample by calling the run_train() and run_evaluation() functions, respectively.

> from shapr import run_train
> run_train()
> run_evaluation()

Dataset

Please find the datasets used for this publication on Zenodo: https://doi.org/10.5281/zenodo.7031924

Folder structure

SHAPR expects the data in the following folder structure (see sample). With corresponding files having the same name. 2D microscopy images (64x64px) should be contained in the images folder, 2D segmentations (64x64px) in the mask folder and the 3D segmentation (64x64x64px) in the obj folder.

path
├── image                 
│   ├── 000003-num1.png
│   │── 000004-num9.png
│   │── 000006-num1.png
│   │── .
│   │── .
│   │── .
│   │── 059994-num1.png     
│
└── mask               
│   ├── 000003-num1.png
│   │── 000004-num9.png
│   │── 000006-num1.png
│   │── .
│   │── .
│   │── .
│   │── 059994-num1.png    
│
└── obj     
│   ├── 000003-num1.png
│   │── 000004-num9.png
│   │── 000006-num1.png
│   │── .
│   │── .
│   │── .
│   │── 059994-num1.png    

Additional analyses

Next to the SHAPR model, we also provide additional scripts that permit zooming into various aspects of SHAPR and persistent homology. All of these scripts are to be found in the scripts subdirectory.

Evaluation of results

To create the evaluation plots with all errors, as shown in the paper, run the evaluation.py script:

$ python -m scripts.evaluation evaluation/red-blood-cell.json

You can optionally specify the -q option in order to skip evaluation metrics that are computationally more complex to compute.

Gallery of results

After running SHAPR, you can use the make_gallery.py script to create a small gallery of all outputs. This is great for getting an overview of how the model performed:

# Note that we are using a sample directory that contains ground truth
# objects. In a realistic scenario, you would rnu this command for the
# outputs of your own model.
$ python -m scripts.make_gallery ../docs/sample/obj 

Interpolation errors

To run the analysis of the interpolation (i.e. downsampling) quality as described in the supplementary materials, run the following command from the shapr directory:

$ python -m scripts.analyse_interpolation -s 8 -p config/small-0D.json

This will generate the interpolation/downsampling errors for a volume of side length s = 8.

Persistent homology of images

To calculate persistent homology of input images, you can use the following script:

$ python -m scripts.calculate_persistence_diagrams ../docs/sample/obj/*.tif

Contributing

We are happy about any contributions. For any suggested changes, please send a pull request to the develop branch.