Comparative study of U-Net, Attention U-Net and UC-TransNet for medical tumor segmentation
- Enrica Bongiovanni
- Laura Prendin
- Arianna Stropeni
This project explores how the standard U-Net architecture can be improved for medical image tumor segmentation by incorporating attention mechanisms and transformer-based modules.
We perform a comparative analysis of three architectures:
- U-Net (baseline)
- Attention U-Net
- UC-TransNet
The models are evaluated on three different medical imaging datasets, each with a different modality:
- Breast ultrasound
- Skin dermoscopy
- Brain MRI
Despite similar quantitative metrics, attention-based models show visually more precise segmentation masks, especially in challenging scenarios.
| Model | Description |
|---|---|
| U-Net | Baseline convolutional architecture for biomedical segmentation |
| Attention U-Net | U-Net with attention gates to focus on relevant regions |
| UC-TransNet | U-Net enhanced with channel-wise transformers to reduce semantic gaps |
The project includes the following files:
- Utils: it defines the functions required to create the datasets and the metrics used
- Network layers: it defines the functions for the custom layers used in the networks
- Base U-Net, Attention U-Net, UC-TransNet: they contain the architecture definition and the training on the three datasets of the different models
- Evaluation: it includes the evaluation of the best models on the test set and the visualization of the results
- Best Models: the folder includes files .h5 that store the best models for each dataset and each architecture
To run the code, download the datasets from the following links, unzip them and store them in a "Datasets" folder in the repository:
- Breast cancer dataset
- Skin cancer dataset: download training, validation and test data and the associated ground truths only for task 1
- Brain cancer dataset: rename the unzipped file "Brain_tumor_dataset"