Skip to content

Abdomen 3D Segmentation Using UNETR: Tool for segmenting abdominal organs using the UNETR model. Combines Transformers with CNNs for precise 3D segmentation. Dive into medical imaging! πŸ₯βœ¨πŸ”

License

Notifications You must be signed in to change notification settings

KaranAnchan/Unetr_3D_Abdomen_Segmentation

Repository files navigation

πŸ₯ Abdomen 3D Segmentation Using UNETR πŸš€

This project uses the UNETR model for precise 3D segmentation of abdominal organs, combining Transformer-based encoders with CNN decoders for good accuracy.

UNETR

πŸ“ Project Structure

UNETR_3D_Abdomen_Segmentation/
β”‚
β”œβ”€β”€ data/
β”‚   β”œβ”€β”€ imagesTr/
β”‚   β”œβ”€β”€ imagesTs/
β”‚   β”œβ”€β”€ labelsTr/
β”‚   └── dataset_0.json
β”‚
β”œβ”€β”€ src/
β”‚   β”œβ”€β”€ __init__.py
β”‚   β”œβ”€β”€ data_loader.py
β”‚   β”œβ”€β”€ model.py
β”‚   β”œβ”€β”€ train.py
β”‚   β”œβ”€β”€ evaluate.py
β”‚   β”œβ”€β”€ utils.py
β”‚   └── visualize.py
β”‚
β”œβ”€β”€ requirements.txt
β”œβ”€β”€ README.md
└── .gitignore
  • data/: Contains training and testing images and labels in .nii.gz format.
  • src/: Source code for the project, organized into modules:
    • data_loader.py: Data loading and preprocessing functions.
    • model.py: Definition of the UNETR model.
    • train.py: Training script for model training.
    • evaluate.py: Script for model evaluation.
    • utils.py: Utility functions for model management.
    • visualize.py: Tools for visualizing predictions.
  • requirements.txt: Python package dependencies.
  • README.md: Documentation of the project.
  • .gitignore: Files and directories to be ignored by Git.

πŸš€ Getting Started

Follow these instructions to set up and run the project on your local machine.

πŸ“‹ Prerequisites

Ensure you have the following installed:

  • Python 3.7+
  • Pip (Python package manager)

πŸ”§ Installation

  1. Clone the repository:

    git clone https://github.com/yourusername/Abdomen-3D-Segmentation-UNETR.git
    cd Abdomen-3D-Segmentation-UNETR
  2. Install dependencies:

    pip install -r requirements.txt

πŸƒβ€β™‚οΈ Usage

Training the Model

To start training the UNETR model:

python src/train.py --data_dir data/ --epochs 100 --batch_size 2
  • --data_dir: Path to the data directory.
  • --epochs: Number of training epochs (default is 100).
  • --batch_size: Batch size for training (default is 2).

Evaluating the Model

Evaluate the model on the validation dataset:

python src/evaluate.py --data_dir data/ --model_path models/model_epoch_100.pth
  • --model_path: Path to the saved model checkpoint.

πŸ“Š Results

The model achieves an impressive Dice score of 0.8027 on the validation set, indicating high accuracy in segmenting abdominal organs.

LOSS And METRIC Visualization

ITK SNAP Visualization

🎨 Visualization

To visualize the model predictions:

from src.visualize import visualize
import nibabel as nib

# Load sample data
image = nib.load('path_to_image.nii.gz').get_fdata()
label = nib.load('path_to_label.nii.gz').get_fdata()
prediction = nib.load('path_to_prediction.nii.gz').get_fdata()

# Visualize a specific slice
visualize(image, label, prediction, slice_idx=60)

This script displays the image, ground truth label, and model prediction side by side for a specified slice.

SLICE VISUALIZATION

πŸ“š Documentation

  • data_loader.py: Handles loading and preprocessing of training and validation data.
  • model.py: Contains the architecture for the UNETR model.
  • train.py: Manages the training process, including loss calculation and model checkpointing.
  • evaluate.py: Evaluates the model's performance on the validation dataset using the Dice metric.
  • utils.py: Utility functions for saving and loading model checkpoints.
  • visualize.py: Provides visualization functions for model predictions.

πŸ› οΈ Built With

  • MONAI - Medical Open Network for AI
  • PyTorch - Deep learning framework
  • NumPy - Numerical computations
  • Matplotlib - Plotting and visualization

πŸ“„ License

This project is licensed under the MIT License - see the LICENSE file for details.

πŸŽ‰ Acknowledgments

  • The developers of MONAI for providing excellent medical imaging tools.
  • The creators of the UNETR architecture for their innovative approach to medical image segmentation.

About

Abdomen 3D Segmentation Using UNETR: Tool for segmenting abdominal organs using the UNETR model. Combines Transformers with CNNs for precise 3D segmentation. Dive into medical imaging! πŸ₯βœ¨πŸ”

Topics

Resources

License

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published

Languages