Multiple Instance Learning for Histopathology: An Introduction to torchmil¶

Author: Francisco Miguel Castro Macías (see my webpage here; contact me at francastro8b@gmail.com).

Summary: In this tutorial, we will learn how to use the torchmil library to train a Multiple Instance Learning (MIL) model to detect breast cancer metastases from Whole Slide Images (WSIs). We will use the CAMELYON16 dataset, which is a widely used benchmark for this task. We will prepare the data using the TRIDENT library. Then, we will train a simple attention-based MIL model and evaluate its performance. Finally, we will visualize the results and localize the metastases in the WSIs.

Presentation and code: You can find the presentation associated with this tutorial here. The code for this tutorial is available in this GitHub repository.

Environment: It is recommended to create two separate conda environments for this tutorial. One for the patch extraction process using TRIDENT and another one for training the MIL model using torchmil. The following commands will create both environments and install the necessary dependencies.

# TRIDENT environment
conda create -n trident python=3.10 -y
conda activate trident
cd some_directory_where_you_want_to_clone_trident
git clone https://github.com/mahmoodlab/trident.git
cd trident
pip install -e .

# torchmil environment
conda create -n torchmil python=3.11 -y
conda activate torchmil
pip install torchmil openslide-bin openslide-python opencv-python scikit-learn matplotlib

Data: You will need to download the original CAMELYON16 dataset from here.

Let’s start by importing the necessary libraries, setting up the device for training, and defining the paths to the data.

1. Preparing and Understanding The Data¶

Whole Slide Images (WSIs). We will be working with Whole Slide Images (WSIs). These are gigapixel-sized images that are used in digital pathology. They are typically obtained by scanning glass slides that contain tissue samples. WSIs are used for various applications, including cancer diagnosis, prognosis, and research. They are typically stored in a multi-resolution format, which allows for efficient storage and retrieval of different levels of detail. We will be working with the CAMELYON16 dataset, which contains WSIs of lymph node sections from breast cancer patients. The WSIs in the CAMELYON16 dataset are stored in the .tif format.

Segmentation and patch extraction. Due to the large size of WSIs, it is not feasible to process them directly with deep learning models. Instead, we need to segment the tissue regions and extract smaller patches from these regions. This process is known as patch extraction. TRIDENT is arguably the most popular tool to accomplish this. I recommend taking a look at the TRIDENT repository for more details on how to use it.

The following commands will run the patch extraction process for all the WSIs in the ./wsis directory and save the processed data in the ./trident_processed directory. We will be using the CONCH model as patch encoder (we will give more details on what a patch encoder is soon). The patches will be extracted at a magnification level of 20x and will have a size of $512 \times 512$ pixels (which is equivalent to a patch size of $1024 \times 1.024$ at 10x magnification).

cd some_directory_where_you_cloned_trident
python run_batch_of_slides.py --task all --wsi_dir /data/datasets/CAMELYON16/original --job_dir /data/datasets/CAMELYON16/trident_processed  --patch_encoder conch_v15 --mag 20 --patch_size 512

This command will create the following directory structure in the ./trident_processed directory:

trident_processed/
└── 20x_512px_0px_overlap/
    ├── features_conch_v15/
    |   ├──wsi1.h5
    |   ├──wsi2.h5
    |   └──...
    └── patches/
       ├──wsi1_patches.h5
       ├──wsi2_patches.h5
       └──...

The features_conch_v15 directory contains the feature vectors for each patch, which are obtained by passing the patches through the CONCH model. The patches directory contains the coordinates of each patch in the original WSI. We will use these coordinates to assign labels to each patch based on the tumor annotations provided in the CAMELYON16 dataset.

Extracting patch labels. The patch extraction process will give us the coordinates of each patch and its corresponding feature vector (obtained from the CONCH model). For CAMELYON16, we have annotations that indicate the tumor regions in each WSI. This means that we can assign a label to each patch based on whether it falls within a tumor region or not. Although these patch labels will not be used for training the MIL model, they can be useful for evaluating the performance of the model at the patch level (e.g., by computing patch-level accuracy, precision, recall, etc.).

To assign a label to each patch, we will check if the proportion of tumor pixels in the patch is greater than a certain threshold (e.g., 50%). If it is, we will assign a label of 1 (indicating that the patch is positive for tumor), otherwise we will assign a label of 0 (indicating that the patch is negative for tumor). We provide a script that accomplishes this:

python extract_labels_camelyon16.py \ 
    --csv_path /data/datasets/CAMELYON16/original/wsi_labels.csv \ 
    --masks_dir /data/datasets/CAMELYON16/original/masks \
    --coords_dir /home/fran/data/datasets/CAMELYON16/trident_processed/20x_512px_0px_overlap/patches \
    --save_dir /home/fran/data/datasets/CAMELYON16/trident_processed/20x_512px_0px_overlap/patch_labels/ \
    --patch_size 1024 --patch_level 0

This will create a new directory called patch_labels that contains a .h5 file for each WSI with the labels for each patch.

At this point, we have prepared all the data that we need. Let’s take a look at some of the patches that we have extracted from a WSI. For this, we read the region determined by the coordinates of each patch.

As you can see, each patch contains a portion of the tissue. We can visualize the whole WSIs by placing each patch in its corresponding location in a blank canvas. Let’s first read all the patches for the selected WSI.

Now that we have the patches, we can visualize the WSI using the patches_to_canvas function from the torchmil.visualize module. This function takes the list of patches and their corresponding coordinates and creates a canvas that represents the whole WSI. To visualize the labels of the patches, we use the draw_heatmap_wsi function from the same module. This function takes the canvas, the patch labels, and the patch coordinates and draws a heatmap on top of the canvas, where red indicates positive patches and green indicates negative patches.

2. Creating the dataset¶

Let’s now create a dataset to train our MIL models. For this, we will use the TridentWSIDataset class from the torchmil.datasets module. This dataset class is designed to work with the data processed by TRIDENT and allows us to easily load the features and labels for each WSI. We will create a training set and a test set. The training set will contain 80% of the WSIs and the test set will contain the remaining 20%.

We provide two CSV files with the labels for each WSI in the training and test sets. Each CSV file has two columns: wsi_name and wsi_label. The wsi_name column contains the name of the WSI (without the .tif extension) and the wsi_label column contains the label for that WSI (0 for normal and 1 for tumor). You can modify these CSV files to create your own training and test sets. Just make sure that the wsi_name column matches the names of the WSIs in the trident_processed directory.

As you can see, in torchmil each bag is a TensorDict, and the different keys correspond to different elements of the bag. In this case, each bag has a feature matrix X, the bag label Y, the instance labels y_inst, the coordinates of the patches coords and an adjacency matrix adj. The adjacency matrix is used to define the spatial relationships between instances in the bag. Recall that the instance labels cannot be used during training, they are available only for evaluation purposes.

3. Mini-batching of bags¶

Tipically, the bags in a MIL dataset have different size. This can be a problem when creating mini-batches. To solve this, we use the function collate_fn from the torchmil.data module. This function creates a mini-batch of bags by padding the bags with zeros to the size of the largest bag in the batch. The function also returns a mask tensor that indicates which instances are real and which are padding.

Let’s create the dataloaders and visualize the shapes of a mini-batch.

Each batch is again a TensorDict with an additional key mask that indicates which instances are real and which are padding. As we can see, the bags are padded to the maximum size of the bags in the batch with zeros. The mask tensor indicates which elements are real instances and which are padding. The function collate_fn also pads other tensors, such as the adjacency matrix or the instance coordinates.

4. Training a model in CAMELYON16¶

We have shown how to load the CAMELYON16 dataset for the binary classification task. Now, let us train a MIL model in this dataset!

For this example, we are going to implement the simple ABMIL model. ABMIL is already implemented in torchmil, (see here). However, for educational purposes, we will implement it from scratch. ABMIL is a simple attention-based MIL model that uses a single attention layer to aggregate the instance features into a bag representation. Despite its simplicity, ABMIL has been shown to achieve competitive performance in various MIL benchmarks, including CAMELYON16.

Recall how ABMIL was defined:

\begin{aligned} \mathbf{h}_i & = \operatorname{NeuralNet}(\mathbf{x}i) \ f_i & = \mathbf{w}^\top \tanh(\mathbf{h}i \mathbf{V}) \ s_i & = \frac{\exp(f_i)}{\sum{j} \exp(f_j)} \ \mathbf{z} &= \sum{i} s_i \mathbf{h}_i \ \hat{Y} & = \operatorname{Classifier}(\mathbf{z}) \end{aligned}

We have implemented a method called compute_loss that computes the binary cross-entropy loss for a given batch of bags. We have done this so that we can train the model using the Trainer class from torchmil. This class provides a generic training loop for any MIL model and allows us to easily log the losses and metrics during training. You can check the documentation of the Trainer class for more details on how to use it and the different parameters that it accepts.

For now, let us just keep it simple and perform a simple training using the torch.optim.Adam optimizer. When using the features from a foundation model, most of torchmil models obtain very good results in just a few epochs.

5. Evaluating the Model¶

Let’s evaluate the model. We are going to compute the accuracy and f1-score on the test set. The accuracy is the proportion of correctly classified bags, while the f1-score is the harmonic mean of precision and recall. The f1-score is a good metric for imbalanced datasets. Typically, in MIL datasets, there are more negative bags than positive bags.

Excellent! Our model has reached a very high accuracy and f1-score in only 5 epochs! This shows how simple is to obtain very good results in one of the most famous WSI classification datasets.

6. Localizing the Metastases¶

Now that we have trained a model, we can use it to localize the metastases in the WSI. We will use the attention scores of the instances to localize the metastases. The attention scores indicate the importance of each instance for the bag-level prediction. Patches with high attention scores are more likely to contain tumor tissue, while patches with low attention scores are more likely to contain normal tissue.

Let’s choose the same WSI that we visualized at the beginning of the tutorial and compute the attention scores for each patch.

Now, we visualize the attention scores as a heatmap on top of the WSI. As before, we use the function draw_heatmap_wsi from the torchmil.visualize module. The heatmap shows the predicted scores for each instance, with red indicating high scores and green indicating low scores. You can also customize these colors.

The heatmap shows the predicted scores for each instance, with red indicating high scores and green indicating low scores. Our model is able to detect the metastases in the WSI! Note that the model has learned to detect these regions without any instance-level labels, only using the bag-level labels. This is a clear example of the power of MIL in WSI classification tasks.