Project 7 - Improved 3D U-Net for Prostate MRI Segmentation (Abhya Garg- 48299785)

[abhyagarg22]

Improved 3D U-Net for Prostate MRI Segmentation

This project implements an Improved 3D U-Net architecture to perform multi-class segmentation of prostate MRI scans from the HipMRI Study Open Dataset.
The goal is to automatically identify and delineate anatomical structures (prostate zones, bladder, rectum, etc.) from volumetric MRI data, reducing the need for time-consuming manual annotations.

The model enhances the traditional 3D U-Net by incorporating residual connections, dropout regularization, and a multi-class Dice + Cross Entropy loss, leading to more stable training and better generalization across classes.

Training was conducted on the UQ Rangpur GPU Cluster (A100) for 15 epochs, with consistent convergence and reliable segmentation quality.

---

Key Features

• Architecture: Encoder–decoder 3D U-Net with skip connections and residual blocks
• Loss Function: Weighted Dice + Cross Entropy for multi-class segmentation
• Regularization: Dropout (0.3) and instance normalization
• Output: 6-class segmentation with softmax activation
• Dataset: /home/groups/comp3710/HipMRI_Study_open/ (semantic MRI & labels)

---

Implementation Highlights

• modules.py: Defines the Improved 3D U-Net model
• dataset.py: Loads and normalizes 3D MRI and label volumes
• train.py: Handles GPU training and Dice evaluation per epoch
• predict.py: Generates predicted segmentation maps and saves results
• train_gpu.slurm: SLURM job script for automated GPU execution
• visualization/: Contains final loss and Dice coefficient plots, segmentation visuals

---

Results

Metric	Value
Final Dice Coefficient	0.742
Average Training Loss	0.455
Epochs	15
GPU	NVIDIA A100

The Dice coefficient improved steadily across epochs, reaching 0.74, confirming effective learning and boundary recognition for prostate structures.

---

Conclusion

The Improved 3D U-Net demonstrated robust segmentation capabilities on 3D MRI data.
Residual learning, dropout, and Dice optimization contributed to smoother boundary detection and stable convergence.
With further fine-tuning (e.g., extended epochs or data augmentation), the model can achieve near-clinical accuracy for prostate delineation tasks.

[wangzhaomxy]

<This is an initial inspection, no action is required at this point.>

File Organizing:

• The README file is currently placed in the wrong directory. Please move it to your project folder at
  /PatternAnalysis-2025/recognition/unet_3d_segmentation
  instead of the current root directory:
  /PatternAnalysis-2025/recognition/.
• In addition, please restore the original README file in the root folder /PatternAnalysis-2025/ and /PatternAnalysis-2025/recognition/.
• Kindly make these corrections; otherwise, the merge request will be rejected.

Problem Solving:

• The algorithm solves the problem appropriately.
• Accuracy in testing dataset (Dice): No test results in the report. Just show training and validation results.

Model and functions:

• It correctly uses PyTorch to construct the improved UNet 3D models and functions.
• Good data augmentation.
• Properly split and use the train/validation/test datasets.

Code design: Good.

Code comment and docstring:

• Good code comments
• Good function docstrings
• Minimal header block

Difficulty: Hard.

Additional Comments:

• Good commits
• The README design can be more structured but the report content is comprehensive. Well done!

[abhyagarg22]

I have now fixed the readme files and about the dice test score I have pasted the picture in the readme and you can see it in my report as well. My dice score test result was able to show in the report I submitted on blackboard. For instance, here is my test score again.

…

On Sat, Nov 22, 2025 at 7:10 AM Zhao WANG ***@***.***> wrote: *wangzhaomxy* left a comment (shakes76/PatternAnalysis-2025#285) <#285 (comment)> *<This is an initial inspection, no action is required at this point.>* *File Organizing:* - *The README file is currently placed in the wrong directory. Please move it to your project folder at /PatternAnalysis-2025/recognition/unet_3d_segmentation instead of the current root directory: /PatternAnalysis-2025/recognition/.* - *In addition, please restore the original README file in the root folder /PatternAnalysis-2025/ and /PatternAnalysis-2025/recognition/.* - *Kindly make these corrections; otherwise, the merge request will be rejected.* *Problem Solving:* - The algorithm solves the problem appropriately. - Accuracy in testing dataset (Dice): *No test results in the report. Just show training and validation results.* *Model and functions:* - It correctly uses PyTorch to construct the improved UNet 3D models and functions. - Good data augmentation. - Properly split and use the train/validation/test datasets. *Code design:* Good. *Code comment and docstring:* - Good code comments - Good function docstrings - *Minimal* header block *Difficulty:* Hard. *Additional Comments:* - Good commits - *The README design can be more structured but the report content is comprehensive. Well done!* — Reply to this email directly, view it on GitHub <#285 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/BJCKZBYTYW2X3TS2DZBLBKT3555SVAVCNFSM6AAAAACLLOFGVOVHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMZTKNRUGU4TMNJSHE> . You are receiving this because you authored the thread.Message ID: ***@***.***>

[abhyagarg22]

I even added dice score output log again as Screenshot.png
Please check it out
Thank you

[gayanku]

Marking

Good/OK/Fair Practice (Design/Commenting, TF/Torch Usage)
	Good design and implementation.
	Spacing and comments.
	Header blocks.
Recognition Problem
	Good solution to problem.
	Driver Script present.
	File structure present.
	Good Usage & Demo & Visualisation & Data usage.
	Module present.
	Commenting present.
	No Data leakage found.
	Difficulty : Hard. Hard. ImprovedUnet3D
Commit Log
	Good Meaningful commit messages.
	Good Progressive commits.
Documentation
	Readme :Acceptable.	-2
	Model/technical explanation :Good.
	Description and Comments :Good.
	Markdown used and PDF submitted.
Pull Request
	Successful Pull Request (Working Algorithm Delivered on Time in Correct Branch).
	Feedback action require: Feedback marks possible +2 if the requested changes are made. Fix file locations for merge.	-2
	Request Description is good.
TOTAL		-4

Marked as per the due date and changes after which aren't necessarily allowed to contribute to grade for fairness.
Subject to approval from Shakes

[abhyagarg22]

Hi I already fixed and made the requested changes to fix the file locations. Kindly have a look at them and please reconsider the feedback action require marks.