COMP3710 Project 3: Improved 2D U-Net for HipMRI Prostate Segmentation

[prabhjotsingh1313]

Project Summary

Implementation of Improved 2D U-Net for prostate cancer segmentation on HipMRI dataset using PyTorch.

Key Achievements

• Prostate Dice coefficient: 0.9373 (exceeds 0.75 requirement by 24.9%)
• Mean Dice across all classes: 0.9158
• All anatomical structures accurately segmented
• Complete modular implementation with comprehensive documentation

Problem Solved

Project 3 from project specification: Segment HipMRI Study on Prostate Cancer using 2D Improved U-Net with minimum Dice similarity coefficient of 0.75 on prostate label.

Files Included

• modules.py: U-Net architecture with dilated convolutions and skip connections
• dataset.py: NIfTI data loading with preprocessing pipeline
• train.py: Complete training, validation, and testing pipeline
• predict.py: Inference script with visualization capabilities
• README.md: Comprehensive documentation with results and usage
• requirements.txt: All dependencies with versions
• images/: Training curves, predictions, and overlay visualizations

Model Architecture

• 5-level Improved U-Net with deeper bottleneck
• Dilated convolutions for increased receptive field
• Batch normalization for training stability
• ~7.76M trainable parameters

Results on Test Set

Channel	Dice Coefficient
0 (Background)	0.9952
1	0.9768
2	0.9023
3 (Prostate)	0.9373
4	0.8717
5	0.8113

Mean Dice: 0.9158

Training Configuration

• Framework: PyTorch 1.12+
• Optimizer: Adam (lr=1e-3)
• Loss: Dice Loss
• Batch size: 4
• Image size: 256×256
• Training time: ~2 minutes/epoch on GPU

Testing Instructions

# Install dependencies
pip install -r requirements.txt

# Train model
python train.py --data_path /path/to/HipMRI_2D --epochs 20 --batch_size 8

# Run predictions
python predict.py --data_path /path/to/HipMRI_2D --checkpoint checkpoints/best_model.pth

Documentation

Complete README.md included with:

• Algorithm description and architecture diagram
• Data preprocessing justification
• Dataset split rationale
• Dependencies with versions
• Usage examples
• Quantitative results
• Visualization examples

---

Student: Prabhjot Singh
Student ID: s4884308
Course: COMP3710 Pattern Analysis
Difficulty: Normal (capped at 15/20 marks)
Requirement: Prostate Dice ≥ 0.75 Achieved: 0.9373