Tamir Shor (tamir.shor@campus.technion.ac.il), Chaim Baskin, Alex Bronstein
This repo contains official Pytorch implementation for the paper "Leveraging Latents for Efficient Thermography Classification and Segmentation".
Our algorithm focuses on using an expressive latent space in order to perform efficient, accurate and fully automatic tumor benign/malignant classification and novel 7-region semantic segmentation. This approach replaces previous methods focused on complex feature selection pipelines or hevay, difficult to train neural architectures. Using our potent latent space, we show it suffices to train a relatively small, simple decoder network to achieve accurate results over the two downstream tasks.
Our classification produces SOTA results, while we are the first to solve the 7-region semantic segmentation problem (segmenting left/right breasts, left/right nipples, left/right armpits and neck).
As a first step you must mount your dataset under the "data" folder. In the paper we used the open DMR dataset. To recreate our experiments place the DMR folder under the data folder. If you wish to use your own dataset, you should also place it under the data folder, however you must also add your custom dataset classes to "src/datasets", similar to "src/datasets/dmr_dataset.py".
You can install the conda environment we used with:
conda env create -f environment.yaml conda activate llft
Our algorithm consists of two decoupled parts - training of the encoder to learn the expressive latent space, and training the decoder to achieve the downstream task.
The encoder is trained in an unsupervised fashion, we follow the CUTS contrastive learning-based architecutre proposed in this paper, implemented in this repo. In our experiments described in the paper we compare between training the encoder over grayscale thermographies and RGB-heatmap thermographies. The choice between data types is controlled by the --gray flag. To train the encoder you must run the main_latent.py script as follows:
python main_latent.py --gray <If if set, train over grayscale data > --cfg-path <If path to folder containing config yaml files. ./config by default>
If you are using your own custom dataset, make sure to add your config file pointing to this dataset to the config folder. See example in "config/dmr_cfg.yaml".
For our decoder we use a basic UNet model implemented under the "model" folder. You can add other architectures there if needed. In our experiments in the paper we consider two downstream tasks - fully supervised tumor classification and few-shot region semantic segmentation. For each task we compare between 4 cases using an encoder trained on heatmap/grayscale data, and for each of those using a decoder trained on heatmap/grayscale data. To train for the classification task:
python main_binary.py --gray_encoder <If set, use encoder trained on grayscale data, else heatmap data> --gray_decoder <if set, use decoder trained on grayscale data, else heatmap data> --cfg-path <path to config folder, "./config" by default> --encoder-ckpt-root <path to folder containing the folders dmr_runs and dmr_runs_gray, containing the respective trained encoder checkpoints>
To train the segmentation task:
python main_seg.py ...
The arguments are the same as for the main_binary.py script above. Note that to train the segmentation task, you'll need some segmentations labels under the path "data/DMR/". Our own labels are under this path in this repo ("52_instances_default.json").
- Our encoder model, as well as the encoder training script and several utility scripts used in our repo, are based on the official CUTS implementation.
- We extend our gratitude to ThermoMind for supplying us with the segmentation annotations.
@article{shor2024leveraging,
title={Leveraging Latents for Efficient Thermography Classification and Segmentation},
author={Shor, Tamir and Baskin, Chaim and Bronstein, Alex},
journal={arXiv preprint arXiv:2404.06589},
year={2024}
}