Uncertainty Estimation in Landmark Localization

1. Description

In this example we implement the methods from Uncertainty Estimation for Heatmap-based Landmark Localization. The method is Quantile Binning, which bins landmark predictions by any continuous uncertainty estimation measure. We assign each bin estimated localization error bounds. We can use these bins to filter out the worst predictions, or identify the likely best predictions.

We evaluate how well an uncertainty measure predicts localization error by measuring the Jaccard Index (a similarity measure) between the predicted bins and the ground truth error quantiles. We also evaluate the accuracy of the estimated error bounds. This framework is applicable to any dataset consisting of (Continuous Uncertainty Measure, Continuous Evaluation Metric) tuples.

Fig. 1 depicts the features exemplified in this example. Note that a) and b) are precomputed and provided in tabular form for this example.

Fig. 1: Overview of our general Quantile Binning framework. a) We make a prediction using a heatmap-based landmark localization model, and b) extract a continuous uncertainty measure. c) We learn thresholds to categorize predictions into bins of increasing uncertainty, estimating error bounds for each bin. e) We filter out predictions from high uncertainty bins to improve the proportion of acceptable predictions. d) Finally, we evaluate each uncertainty measure's ability to capture the true error quantiles and the accuracy of the estimated error bounds.

2. Datasets

We provide three tabular datasets containing landmark localization error and uncertainty estimation values: 1) 303 Short Axis View CMR images with 3 landmarks (SA), 422 Four Chamber View CMR images with 3 landmarks (4ch), and 400 Cephalometric Radiology images with 19 landmarks (Cephalometric). The CMR dataset is from the ASPIRE Registry, and the Cephalometric dataset is from an ISBI grand challenge.

For each uncertainty measure we provide tuples of (Continuous Uncertainty Measure, Continuous Localization Error) for each sample in the validation and test set in tabular form. We have split the data into 8 folds and used cross validation to gather validation and test set uncertainty tuples for every sample in the datasets.

In this example, we compare the uncertainty measures:

• Single Maximum Heatmap Activation (S-MHA).
• Ensemble Maximum Heatmap Activation (E-MHA).
• Ensemble Coordinate Prediction Variance (E-CPV).

We compare these measures on landmark predictions from:

• A U-Net model.
• A PHD-Net model.

3. Usage

Clone the repo

git clone https://github.com/Schobs/Qbin

Install the required packages:

conda create --n qbin_env, conda activate qbin_env  (Optional)

pip install -r requirements.txt

Run Quantile Binning for the Four Chamber, Short Axis, and Cephalometric data respectively:

cd code
python main.py --cfg configs/4CH_data.yaml 
python main.py --cfg configs/SA_data.yaml
python main.py --cfg configs/isbi_config.yaml

Tip: To save all figures automatically instead of displaying, set SAVE_FIGURES: True in the .yaml file.

Edit it the above yaml files for additional configuration options.

4. Additional Experimental Results

Find additional experimental results from our paper in the additional results folder.