suhitaghosh10/UATS

A deep semi-supervised method (UATS) for medical segmentation

Uncertainty-Aware Temporal Self-Learning (UATS)

This repository contains the source code of the paper Uncertainty-Aware Temporal Self-Learning (UATS): Semi-Supervised Learning for Segmentation of the Prostate Zones and Beyond, accepted in Artificial Intelligence in Medicine Journal.

Highlights:

• A deep semi-supervised method named uncertainty-aware temporal self-learning (UATS) is proposed for segmentation.
• UATS leverages performance gains from temporal ensembling (Laine and Aila) and uncertainty-guided self-learning.
• UATS surpasses fully supervised performance on prostate zone segmentation and achieves human observer quality.
• Further experiments demonstrate its generalizability on the following benchmark biomedical datasets:
  • Prostate Zones Segmentation
  • ISIC2018 Skin Lesion Segmentation challenge
  • Medical Decathlon Challenge Task 04 (Hippocampus Segmentation)

Training:

1. pip install -r requirements.txt
2. Experiment and dataset-related changes to be made in the file: utility/config.py
3. The UATS model could be trained: python dataset_specific/prostate/expts/prostate_uats_softmax.py