This project implements a U-Net model for the segmentation of prostate images. It is designed to demonstrate the effectiveness of U-Net architectures for medical image analysis tasks, particularly in segmenting complex structures from varied backgrounds.
The dataset used in this project is sourced from the SAML Dataset, which contains a large collection of annotated medical images for semantic segmentation. These images are specifically tailored for machine learning models to learn the segmentation of anatomical structures in medical scans.
The U-Net model is a convolutional neural network originally designed for biomedical image segmentation. The architecture is structured as a U-shaped network to efficiently learn from a small number of images while achieving precise localizations. Its design includes a contracting path to capture context and an expansive path that enables precise localization, making it highly effective for tasks like medical image segmentation.
Here is a visualization of the model's segmentation output compared to the ground truth:
- Main Dataset: SAML Dataset Link
- Processed Dataset: Processed Dataset
main.py: The primary script for training the U-Net model. It integrates the model setup, training loop, and saving the trained model.model.py: Contains the implementation of the U-Net model architecture.dataloader.py: Defines the PyTorchDatasetandDataLoaderfor handling the loading and preprocessing of image data.pre_process.ipynb: A Jupyter notebook used for the initial preprocessing of the dataset.
To run the training script, simply execute the following command:
python main.py- torch==2.0.1
- torchvision==0.15.2
- numpy==1.22.0
- matplotlib==3.7.1