mhdank/puma_model_checkpoints

Medical Image Segmentation with U-Net Architectures

Project Overview

Medical image segmentation project using various U-Net architectures to segment microscopic tissue images from the PUMA dataset. Handles multi-class segmentation for both tissue structures and nuclei types.

Dataset

PUMA Training Dataset: Kaggle Link

Contains:

• Microscopic tissue images
• Coordinate data for segmented regions
• Color-coded mask images
• Limited samples (addressed via augmentation)

Class Definitions

Tissue Classes (6 classes)

[
    (25,204,25): 5,        # tissue_necrosis
    (127, 51, 51): 1,      # tissue_stroma
    (204, 25, 25): 3,      # tissue_tumor
    (178, 178, 178): 2,    # tissue_blood_vessel
    (255, 153, 0): 4,      # tissue_epidermis
    (255, 255, 255): 0     # tissue_white_background
]


### Nuclei Classes (11 classes)

The nuclei classes are represented by distinct labels:

```python
{
    "background",
    "nuclei_apoptosis", 
    "nuclei_endothelium",
    "nuclei_epithelium",
    "nuclei_histiocyte",
    "nuclei_lymphocyte",
    "nuclei_melanophage",
    "nuclei_neutrophil", 
    "nuclei_plasma_cell",
    "nuclei_stroma",
    "nuclei_tumor"
}

Model Architectures

Implemented U-Net Variants

1. Standard U-Net: A basic encoder-decoder architecture that serves as the baseline model.
2. Attention U-Net: Incorporates attention gates in the skip connections to focus on relevant features and suppress irrelevant ones.
3. Sharp U-Net: Enhances the boundary detection of segmented regions, improving the precision of the segmentation.
4. Multi-Resolution U-Net: Uses multi-scale feature processing to capture both local and global information.

Data Processing

Mask Generation

• Coordinate to Binary Mask Conversion: Coordinates from the dataset are converted into binary masks for each tissue and nuclei class.
• Class-Specific Color Coding: Each class (both tissue and nuclei) is represented by its specific RGB color in the segmentation mask.
• RGB Mask Generation: Segmentation masks are generated in RGB format, where each pixel's color corresponds to a specific class.

Data Augmentation

To overcome the limited dataset size, the following augmentation techniques are applied:

• Random Rotations: Rotating images by random angles to simulate various orientations.
• Flips: Horizontal and vertical flips to introduce variability.
• Color Adjustments: Changes to brightness, contrast, and saturation to simulate different lighting conditions.
• Elastic Transformations: Non-linear deformations to simulate tissue variability.
• Scale Variations: Random scaling of images to mimic different magnifications.

Training Details

• Training Methodology: All models are trained using a supervised approach with labeled masks.
• Loss Function: Cross-entropy loss is used for multi-class segmentation.
• Optimizer: Adam optimizer with learning rate scheduling.
• Metrics: Intersection over Union (IoU), Dice coefficient, and pixel-wise accuracy are used to evaluate the performance of the segmentation models.

Results

The models were evaluated on the test set, and the best-performing architecture (based on the validation IoU) was selected for further analysis. The segmentation performance is visualized using various metrics to compare the predicted masks with the ground truth.