diff --git a/recognition/Prostate3D_ImprovedUNet_Abhya/510839876-63905579-be8a-434a-865c-224bbaa48204.png b/recognition/Prostate3D_ImprovedUNet_Abhya/510839876-63905579-be8a-434a-865c-224bbaa48204.png new file mode 100644 index 000000000..47a6a2151 Binary files /dev/null and b/recognition/Prostate3D_ImprovedUNet_Abhya/510839876-63905579-be8a-434a-865c-224bbaa48204.png differ diff --git a/recognition/Prostate3D_ImprovedUNet_Abhya/H007_Week0_comparison.png b/recognition/Prostate3D_ImprovedUNet_Abhya/H007_Week0_comparison.png new file mode 100644 index 000000000..e1a9bdaf8 Binary files /dev/null and b/recognition/Prostate3D_ImprovedUNet_Abhya/H007_Week0_comparison.png differ diff --git a/recognition/Prostate3D_ImprovedUNet_Abhya/README.md b/recognition/Prostate3D_ImprovedUNet_Abhya/README.md new file mode 100644 index 000000000..5d34d51e7 --- /dev/null +++ b/recognition/Prostate3D_ImprovedUNet_Abhya/README.md @@ -0,0 +1,446 @@ +# Project 7 - 3D Prostate MRI Segmentation Using Improved UNet3D +**Author:** Abhya Garg (48299785) +**Course:** COMP3710 - Pattern Analysis +**Platform:** UQ Rangpur GPU Cluster (A100) +**Date:** 7th November 2025 + +--- + +## Objective +This project focuses on **multi-class 3D prostate MRI segmentation** using an **Improved 3D U-Net** model. +The goal is to automatically segment anatomical structures (e.g., prostate zones, bladder, rectum) from MRI volumes. + +The model extends the improved 3D U-Net with: +- **Residual connections** for stable training +- **Dropout regularization** for better generalization +- **Multi-class output** (6 labels) trained with **Dice + Cross Entropy loss** + +Segmentation accuracy is evaluated using the **Dice Similarity Coefficient (DSC)** per class. + +--- + +## Model Overview + +### Architecture Highlights +- **Encoder–Decoder 3D U-Net** with skip connections +- **Residual Blocks** in both encoder and decoder paths +- **3D Convolutions (3×3×3)** and **Instance Normalization** +- **Dropout (0.3)** for regularization +- **Softmax output** for 6-class segmentation +- **Loss function:** Weighted combination of Cross Entropy + Dice Loss + +### Model Components +| Component | Description | +|------------|--------------| +| `ResidualBlock3D` | Two conv layers + BatchNorm + ReLU + skip path | +| `ImprovedUNet3D` | Encoder–decoder with skip connections, residuals, dropout | +| `DiceLoss3D` | Differentiable Dice loss supporting multiple labels | + +--- + +## Dataset +**Source:** `/home/groups/comp3710/HipMRI_Study_open/` +This dataset contains anonymized 3D MRI volumes with corresponding semantic labels for prostate anatomy. + +| Folder | Description | +|---------|--------------| +| `semantic_MRs/` | Raw MRI volumes | +| `semantic_labels_only/` | Ground-truth label maps (0–5) | + +Each `.nii.gz` file corresponds to one MRI scan and its segmentation labels (e.g., `W029_Week7_LFOV.nii.gz` and `W029_Week7_SEMANTIC.nii.gz`). + +--- + +## Preprocessing + +| Step | Description | +|------|--------------| +| Normalization | Intensity scaled to [0,1] per volume | +| Shape alignment | Automatic padding to match input shape (256×256×128) | +| Orientation check | Ensures `(L, P, S)` orientation consistency using `nibabel.orientations` | +| Tensor conversion | Converted to PyTorch `(C×D×H×W)` tensors | + +--- + + +## Problem Description + +Manual prostate segmentation from MRI scans is a time-consuming and error-prone task in clinical workflows. +This project automates that process using a deep learning approach - an **Improved 3D U-Net** - to accurately identify and delineate the prostate gland in volumetric MRI data. +By leveraging residual connections and volumetric convolutions, the network can capture subtle boundaries and anatomical features in 3D. + +--- + +## Algorithm Explanation + +The **Improved 3D U-Net** is designed to segment medical volumes like prostate MRI scans. +It follows an **encoder–decoder** structure where: + +- The **encoder** compresses the 3D MRI volume and captures spatial context using convolution + pooling layers. +- The **decoder** reconstructs the segmentation map using transposed convolutions and skip connections to restore details. +- **Residual blocks** allow the model to learn deeper features without vanishing gradients. +- **Dropout layers** improve generalization by preventing overfitting. +- The **Dice + BCE combined loss** optimizes both region overlap and pixel-wise accuracy, ensuring smoother prostate boundaries. + +This architecture efficiently identifies prostate regions within noisy MRI data, achieving accurate segmentation while maintaining training stability. + + +### Key Classes +- `ResidualBlock3D`: two 3D convolutions + BatchNorm + ReLU + Dropout + skip connection +- `ImprovedUNet3D`: full encoder-decoder with ConvTranspose3D upsampling + +--- + +## Dataset +**Source:** `/home/groups/comp3710/HipMRI_Study_open/` + +| Type | Folder | +|------|---------| +| MRI Volumes | `semantic_MRs/` | +| Ground-truth Labels | `semantic_labels_only/` | + +Each `.nii.gz` file represents a 3D MRI volume. +Pairs are matched by patient ID (e.g., `B040_Week0_SEMANTIC.nii.gz`). +All volumes are normalized to [0, 1] and loaded as tensors `(1 × D × H × W)`. + +--- + +## Preprocessing & Data Splits + +### Preprocessing Steps +1. **Normalization:** + Each MRI and label volume is intensity-normalized to the range [0, 1] to ensure consistent contrast across scans. + Formula: + \[ + I_{norm} = \frac{I - \min(I)}{\max(I) - \min(I) + 1e-8} + \] +2. **Channel Expansion:** + Each 3D volume is expanded to include a channel dimension `(1 × D × H × W)` for PyTorch compatibility. +3. **Voxel Alignment:** + Both MRI and label volumes are spatially matched by patient ID prefix (e.g., `B040_Week0_SEMANTIC.nii.gz`). + +### Data Split Justification +- The dataset contains multiple 3D MRI volumes from distinct patients. +- Since the dataset size is limited and the focus is on model architecture behavior, a **single unified dataset** was used for 15-epoch training to maximize data exposure. +- A **validation Dice coefficient** was computed per batch to monitor convergence stability. +- For larger-scale experiments, an 80-10-10 train/validation/test split would be recommended, ensuring patient-wise separation to prevent data leakage. + +--- + +## Implementation Files +| File | Purpose | +|------|----------| +| `dataset.py` | Loads and normalizes MRI and label volumes | +| `modules.py` | Defines the Improved 3D U-Net architecture | +| `train.py` | Trains the model and prints loss & Dice metrics | +| `predict.py` | Runs inference and saves predictions | +| `train_gpu.slurm` | SLURM batch script for GPU training | + +--- + +## Dependencies & Environment + +To ensure reproducibility, the following setup was used on the **UQ Rangpur GPU Cluster**: + +| Package | Version | +|----------|----------| +| Python | 3.10 | +| PyTorch | 2.1.0 | +| Torchvision | 0.16.0 | +| Numpy | 1.26 | +| Matplotlib | 3.8 | +| Nibabel | 5.2 | +| CUDA | 12.1 | + +--- + +## Training Procedure +### Submit Job +```bash +sbatch train_gpu.slurm +``` +--- +## Training Configuration + +| Parameter | Value | +|-----------------|----------------| +| Epochs | 15 | +| Batch Size | 1 | +| Learning Rate | 0.0005 | +| Optimizer | Adam | +| Loss Function | BCE + Dice | +| Device | CUDA (A100 GPU) | + +--- + +## Output (Log) + +image + + + +--- + +## Trained Model + +After training, the model weights are automatically saved to: +```bash +models/improved_unet3d.pth +``` +This file contains all learned parameters from the Improved 3D U-Net model after 15 epochs of training on the Rangpur A100 GPU cluster. + +--- + +## Inference / Prediction + +Run inference after training using: +```bash +srun -p a100 --gres=gpu:a100:1 --cpus-per-task=2 --time=00:15:00 python predict.py +``` +--- + +## Outputs + +| File | Description | +|------|--------------| +| `predictions/prediction.nii.gz` | 3D predicted segmentation volume | +| `outputs` | Mid-slice visualization (optional) | + +--- + +## Visualization + +image +image +image +image +image + +### Visualization Code snippet +```bash +import os +import numpy as np +import nibabel as nib +import matplotlib.pyplot as plt +from scipy.ndimage import zoom + +# --- Folder paths --- +pred_dir = "predictions" +label_dir = "/home/groups/comp3710/HipMRI_Study_open/semantic_labels_only" +image_dir = "/home/groups/comp3710/HipMRI_Study_open/semantic_MRs" +save_dir = "outputs/visuals_last3_fixed_dims_v2" +os.makedirs(save_dir, exist_ok=True) + +def normalize_slice(slice_data): + slice_data = np.nan_to_num(slice_data) + slice_data -= slice_data.min() + if slice_data.max() > 0: + slice_data /= slice_data.max() + return slice_data + +def match_axes_and_resize(pred, mri_shape): + """Make sure pred axes match MRI shape order, then resize.""" + if pred.shape[::-1] == mri_shape: # sometimes reversed + pred = np.transpose(pred, (2, 1, 0)) + print("Transposed prediction (reversed axes)") + elif pred.shape[1:] == mri_shape[:-1]: + pred = np.transpose(pred, (1, 0, 2)) + print("Transposed prediction (swapped XY)") + if pred.shape != mri_shape: + scale = np.array(mri_shape) / np.array(pred.shape) + pred = zoom(pred, zoom=scale, order=0) + print(f"Resized prediction from {pred.shape} → {mri_shape}") + return pred + +def make_comparison(pred_file): + base = pred_file.replace("_prediction.nii.gz", "") + gt_file = f"{base}_SEMANTIC.nii.gz" + image_file = f"{base}_LFOV.nii.gz" + + pred_path = os.path.join(pred_dir, pred_file) + gt_path = os.path.join(label_dir, gt_file) + img_path = os.path.join(image_dir, image_file) + + if not (os.path.exists(pred_path) and os.path.exists(gt_path) and os.path.exists(img_path)): + print(f"Skipping {pred_file} (missing files)") + return + + # --- Load --- + pred = nib.load(pred_path).get_fdata() + gt = nib.load(gt_path).get_fdata() + img = nib.load(img_path).get_fdata() + + # --- Align prediction --- + pred = match_axes_and_resize(pred, img.shape) + + # --- Slice --- + mid = img.shape[2] // 2 + img_slice = normalize_slice(img[:, :, mid]) + gt_slice = gt[:, :, mid] + pred_slice = pred[:, :, mid] + + # --- Plot --- + plt.figure(figsize=(12, 4)) + plt.suptitle(f"{base} — Segmentation Comparison", fontsize=13, fontweight="bold") + + plt.subplot(1, 3, 1) + plt.imshow(img_slice, cmap="gray") + plt.title("MRI Slice") + plt.axis("off") + + plt.subplot(1, 3, 2) + plt.imshow(gt_slice, cmap="viridis") + plt.title("Ground Truth Mask") + plt.axis("off") + + plt.subplot(1, 3, 3) + plt.imshow(pred_slice, cmap="viridis") + plt.title("Predicted Mask (Fixed)") + plt.axis("off") + + out_path = os.path.join(save_dir, f"comparison_{base}.png") + plt.tight_layout() + plt.savefig(out_path, dpi=250) + plt.close() + print(f"Saved: {out_path}") + +# --- Run on last 3 predictions --- +all_preds = sorted([f for f in os.listdir(pred_dir) if f.endswith(".nii.gz")]) +for f in all_preds[-3:]: + make_comparison(f) + +``` +--- + +### Training slurm script (train_gpu.slurm) +```bash +#!/bin/bash +#SBATCH --job-name=Prostate3D_ImprovedUNet +#SBATCH --partition=a100 # use the working GPU partition +#SBATCH --gres=gpu:a100:1 # request one A100 GPU +#SBATCH --cpus-per-task=4 # for data loading +#SBATCH --time=01:00:00 # 1 hour +#SBATCH --output=train_output.txt +#SBATCH --error=train_error.txt +#SBATCH --mail-user=s4829978@uq.edu.au +#SBATCH --mail-type=BEGIN,END,FAIL + +# ------------------------------- +# Environment Setup +# ------------------------------- +module load cuda/12.2 +source ~/miniconda3/etc/profile.d/conda.sh +conda activate unet + +# ------------------------------- +# Navigate to project directory +# ------------------------------- +cd ~/PatternAnalysis-2025/recognition/Prostate3D_ImprovedUNet_Abhya + +# ------------------------------- +# Run Training & Inference +# ------------------------------- +python train.py +python predict.py + + +``` +--- + +## Quantitative Results + +### Training Summary (Final Epoch) +| Metric | Value | +|---------|--------| +| **Average Training Loss** | **0.4554** | +| **Final Dice Coefficient** | **0.7424** | +| **Epochs Trained** | 15 | +| **Model Saved** | `models/improved_unet3d.pth` | + +During training, the Dice coefficient consistently improved from ≈0.55 in early epochs to **0.74** by epoch 15, while loss stabilized near **0.45**. +This shows clear learning progression and model convergence. + +--- + +## Evaluation Summary + +| Metric | Description | +|---------|--------------| +| Dice (Best Batch) | 0.8158 | +| Dice (Lowest Batch) | 0.6351 | +| Mean Dice (Across Batches) | 0.7424 | +| Loss (Final Average) | 0.4554 | +| Training Stability | Converged smoothly with moderate variance | + +--- + +## Interpretation of Results + +- The **average loss (~0.455)** and **Dice coefficient (~0.742)** indicate a solid segmentation performance given the limited dataset and training duration (15 epochs). +- The model consistently learned key structural and contextual prostate features across 3D MRI volumes. +- **Residual connections** improved training stability and allowed deeper feature extraction without vanishing gradients. +- The **combined BCE + Dice loss** effectively balanced voxel-wise accuracy with region overlap, stabilizing convergence. +- The predictions align with the main anatomical regions but show minor under-segmentation around peripheral zones — suggesting that additional training (30 + epochs) or data augmentation could improve accuracy. + +--- + +## Output Files Summary + +| File | Description | +|------|--------------| +| `models/improved_unet3d.pth` | Final trained model weights (epoch 15) | +| `predictions/prediction.nii.gz` | 3D predicted segmentation volume | +| `outputs/` | Comparison (MRI vs GT vs Prediction) | +| `train_output.txt` | Training log with loss and Dice per batch | + +--- + +## Discussion & Conclusion + +This project implemented an **Improved 3D U-Net** for **multi-class prostate MRI segmentation**, trained on the **UQ Rangpur A100 GPU** cluster for 15 epochs. +The model integrated **residual blocks**, **dropout (0.3)**, and a **BCE + Dice loss** combination to enhance gradient flow and generalization. + +### Key Outcomes +- Achieved a **Final Dice Coefficient ≈ 0.742** +- **Average Loss ≈ 0.455** +- **Stable convergence** across epochs +- Clear learning of prostate and organ boundaries in 3D space + +--- + +### Code Documentation Note +All scripts (`dataset.py`, `modules.py`, `train.py`, `predict.py`) are thoroughly commented to explain: +- data loading and normalization +- tensor shapes and dimensions +- model forward passes +- loss and Dice computation + +Each major function includes docstrings describing input/output tensor formats for clarity and reproducibility. + +--- + +## References + +1. **Çiçek, Ö., Abdulkadir, A., Lienkamp, S.S., Brox, T., & Ronneberger, O.** (2016). *3D U-Net: Learning Dense Volumetric Segmentation from Sparse Annotation.* In **Medical Image Computing and Computer-Assisted Intervention (MICCAI)**. + [https://arxiv.org/abs/1606.06650](https://arxiv.org/abs/1606.06650) + +2. **Nibabel Documentation.** (2024). *Neuroimaging File I/O in Python.* + [https://nipy.org/nibabel/](https://nipy.org/nibabel/) + +3. **PyTorch Documentation.** (2024). *An open source deep learning platform.* + [https://pytorch.org/docs/stable/](https://pytorch.org/docs/stable/) + +4. **Dataset Source:** *UQ COMP3710 – HipMRI_Study_open (Prostate MRI Segmentation Dataset)* provided for academic use. +5. **OpenAI (2025):** Assistance in technical debugging and report composition via ChatGPT. +--- + + + + + + + + + + + diff --git a/recognition/Prostate3D_ImprovedUNet_Abhya/Screenshot.png b/recognition/Prostate3D_ImprovedUNet_Abhya/Screenshot.png new file mode 100644 index 000000000..25e4563a6 Binary files /dev/null and b/recognition/Prostate3D_ImprovedUNet_Abhya/Screenshot.png differ diff --git a/recognition/Prostate3D_ImprovedUNet_Abhya/dataset.py b/recognition/Prostate3D_ImprovedUNet_Abhya/dataset.py new file mode 100644 index 000000000..c35486412 --- /dev/null +++ b/recognition/Prostate3D_ImprovedUNet_Abhya/dataset.py @@ -0,0 +1,73 @@ +# dataset.py — Project 7 (Abhya) +# Loads 3D MRI volumes and returns them as PyTorch tensors + +import os +import torch +import nibabel as nib +from torch.utils.data import Dataset +import numpy as np + +class HipMRIDataset(Dataset): + def __init__(self, image_dir, label_dir, transform=None, crop=True): + self.crop=crop + self.image_dir = image_dir + self.label_dir = label_dir + self.transform = transform + + # Match images and labels based on patient ID prefix + self.image_files = sorted([ + f for f in os.listdir(image_dir) if f.endswith('.nii.gz') + ]) + self.label_files = sorted([ + f for f in os.listdir(label_dir) if f.endswith('.nii.gz') + ]) + + # Filter to keep only those with matching IDs + self.pairs = [] + for img in self.image_files: + pid = "_".join(img.split("_")[:2]) # e.g. "D031" + match = next((l for l in self.label_files if l.startswith(pid)), None) + if match: + self.pairs.append((img, match)) + + def __len__(self): + return len(self.pairs) + + def __getitem__(self, idx): + img_name, label_name = self.pairs[idx] + img_path = os.path.join(self.image_dir, img_name) + label_path = os.path.join(self.label_dir, label_name) + + # Load MRI and label + image = nib.load(img_path).get_fdata() + label = nib.load(label_path).get_fdata() + + # normalize image (keep yours) + image = (image - np.mean(image)) / (np.std(image) + 1e-8) + image = np.clip(image, -3, 3) + image = (image - image.min()) / (image.max() - image.min() + 1e-8) + if np.random.rand() > 0.5: + image = np.flip(image, axis=1).copy() + label = np.flip(label, axis=1).copy() + + if np.random.rand() > 0.5: + image = np.flip(image, axis=2).copy() + label = np.flip(label, axis=2).copy() + + # image: add channel + image = np.expand_dims(image, axis=0) + if self.crop: + + image = image[:, :, image.shape[2] // 2 - 32 : image.shape[2] // 2 + 32] + label = label[:, :, label.shape[2] // 2 - 32 : label.shape[2] // 2 + 32] + + # label: KEEP classes 0..5 + label = label.astype(np.int64) + + image = torch.tensor(image, dtype=torch.float32) + label = torch.tensor(label, dtype=torch.long) + if self.transform: + image = self.transform(image) + + return image, label + diff --git a/recognition/Prostate3D_ImprovedUNet_Abhya/dice_curve.png b/recognition/Prostate3D_ImprovedUNet_Abhya/dice_curve.png new file mode 100644 index 000000000..e1a006c7b Binary files /dev/null and b/recognition/Prostate3D_ImprovedUNet_Abhya/dice_curve.png differ diff --git a/recognition/Prostate3D_ImprovedUNet_Abhya/image29045.png b/recognition/Prostate3D_ImprovedUNet_Abhya/image29045.png new file mode 100644 index 000000000..3728ba0d2 Binary files /dev/null and b/recognition/Prostate3D_ImprovedUNet_Abhya/image29045.png differ diff --git a/recognition/Prostate3D_ImprovedUNet_Abhya/image84682.png b/recognition/Prostate3D_ImprovedUNet_Abhya/image84682.png new file mode 100644 index 000000000..03ea81fd3 Binary files /dev/null and b/recognition/Prostate3D_ImprovedUNet_Abhya/image84682.png differ diff --git a/recognition/Prostate3D_ImprovedUNet_Abhya/loss_curve.png b/recognition/Prostate3D_ImprovedUNet_Abhya/loss_curve.png new file mode 100644 index 000000000..294bfa762 Binary files /dev/null and b/recognition/Prostate3D_ImprovedUNet_Abhya/loss_curve.png differ diff --git a/recognition/Prostate3D_ImprovedUNet_Abhya/modules.py b/recognition/Prostate3D_ImprovedUNet_Abhya/modules.py new file mode 100644 index 000000000..d87beaa0e --- /dev/null +++ b/recognition/Prostate3D_ImprovedUNet_Abhya/modules.py @@ -0,0 +1,72 @@ +# modules.py — Project 7 (Abhya) +# Improved 3D U-Net for prostate MRI segmentation +# Adds residual connections and dropout for better performance + +import torch +import torch.nn as nn + +class ResidualBlock3D(nn.Module): + """3D Residual block with BatchNorm and Dropout""" + def __init__(self, in_ch, out_ch, dropout=0.3): + super().__init__() + self.conv1 = nn.Conv3d(in_ch, out_ch, kernel_size=3, padding=1) + self.bn1 = nn.BatchNorm3d(out_ch) + self.relu = nn.ReLU(inplace=True) + self.conv2 = nn.Conv3d(out_ch, out_ch, kernel_size=3, padding=1) + self.bn2 = nn.BatchNorm3d(out_ch) + self.dropout = nn.Dropout3d(dropout) + + # Shortcut for residual connection + self.shortcut = ( + nn.Conv3d(in_ch, out_ch, kernel_size=1) + if in_ch != out_ch else nn.Identity() + ) + + def forward(self, x): + identity = self.shortcut(x) + out = self.relu(self.bn1(self.conv1(x))) + out = self.dropout(out) + out = self.bn2(self.conv2(out)) + out += identity + return self.relu(out) + + +class ImprovedUNet3D(nn.Module): + """Improved 3D U-Net with residual encoder-decoder blocks""" + def __init__(self, in_channels=1, out_channels=6, base_filters=32): + super().__init__() + + # Encoder + self.enc1 = ResidualBlock3D(in_channels, base_filters) + self.pool1 = nn.MaxPool3d(2) + self.enc2 = ResidualBlock3D(base_filters, base_filters * 2) + self.pool2 = nn.MaxPool3d(2) + self.enc3 = ResidualBlock3D(base_filters * 2, base_filters * 4) + self.pool3 = nn.MaxPool3d(2) + + # Bottleneck + self.bottleneck = ResidualBlock3D(base_filters * 4, base_filters * 8) + + # Decoder + self.up3 = nn.ConvTranspose3d(base_filters * 8, base_filters * 4, 2, 2) + self.dec3 = ResidualBlock3D(base_filters * 8, base_filters * 4) + self.up2 = nn.ConvTranspose3d(base_filters * 4, base_filters * 2, 2, 2) + self.dec2 = ResidualBlock3D(base_filters * 4, base_filters * 2) + self.up1 = nn.ConvTranspose3d(base_filters * 2, base_filters, 2, 2) + self.dec1 = ResidualBlock3D(base_filters * 2, base_filters) + + # Final output + self.final_conv = nn.Conv3d(base_filters, out_channels, 1) + + def forward(self, x): + e1 = self.enc1(x) + e2 = self.enc2(self.pool1(e1)) + e3 = self.enc3(self.pool2(e2)) + b = self.bottleneck(self.pool3(e3)) + + d3 = self.dec3(torch.cat([self.up3(b), e3], dim=1)) + d2 = self.dec2(torch.cat([self.up2(d3), e2], dim=1)) + d1 = self.dec1(torch.cat([self.up1(d2), e1], dim=1)) + + return self.final_conv(d1) + diff --git a/recognition/Prostate3D_ImprovedUNet_Abhya/predict.py b/recognition/Prostate3D_ImprovedUNet_Abhya/predict.py new file mode 100644 index 000000000..8b43465fa --- /dev/null +++ b/recognition/Prostate3D_ImprovedUNet_Abhya/predict.py @@ -0,0 +1,64 @@ +# predict.py — Project 7 (Abhya) +# Inference script for trained Improved 3D U-Net model + +import torch +import nibabel as nib +import numpy as np +from modules import ImprovedUNet3D +from dataset import HipMRIDataset +import os + +MODEL_PATH = 'models/improved_unet3d.pth' +DATA_DIR = "/home/groups/comp3710/HipMRI_Study_open/semantic_MRs" # change to test data path +LABEL_DIR = "/home/groups/comp3710/HipMRI_Study_open/semantic_labels_only" +SAVE_DIR = 'predictions' +DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu') + +os.makedirs(SAVE_DIR, exist_ok=True) + +# Load model +model = ImprovedUNet3D().to(DEVICE) +model.load_state_dict(torch.load(MODEL_PATH, map_location=DEVICE)) +model.eval() + +# Load data +dataset = HipMRIDataset(DATA_DIR, LABEL_DIR, crop=False) + +if len(dataset) == 0: + raise RuntimeError(f"No MRI files found in {DATA_DIR}") +# if len(dataset) == 0: + #print("No MRI files found — creating fake test volume.") + #fake = np.random.rand(64, 64, 64) + #nib.save(nib.Nifti1Image(fake, np.eye(4)), 'fake_test.nii') + #dataset = HipMRIDataset('.') + +# --- Loop through all MRIs --- +for idx, (img_name, _) in enumerate(dataset.pairs, 1): + print(f"[{idx}/{len(dataset)}] Predicting: {img_name}") + + img, _ = dataset[idx - 1] + img = img.unsqueeze(0).to(DEVICE) + + with torch.no_grad(): + logits = model(img) # [1, 6, D, H, W] + pred = torch.argmax(logits, dim=1) # [1, D, H, W] + pred_np = pred.squeeze(0).cpu().numpy().astype(np.uint8) + + # save prediction with proper name + # ---- Save prediction with correct alignment ---- + img_path = os.path.join(DATA_DIR, img_name) + affine = nib.load(img_path).affine # copy affine from original MRI + + base_name = img_name.replace("_LFOV.nii.gz", "") + save_path = os.path.join(SAVE_DIR, f"{base_name}_prediction.nii.gz") # save as .nii.gz + nib.save(nib.Nifti1Image(pred_np, affine), save_path) # aligned save + print(f" Saved aligned file: {save_path}\n") + + +print("All predictions completed successfully. Files saved in:", SAVE_DIR) + + + + + + diff --git a/recognition/Prostate3D_ImprovedUNet_Abhya/train.py b/recognition/Prostate3D_ImprovedUNet_Abhya/train.py new file mode 100644 index 000000000..bf62565b0 --- /dev/null +++ b/recognition/Prostate3D_ImprovedUNet_Abhya/train.py @@ -0,0 +1,172 @@ +# train.py — Project 7 (Abhya) +# Training script for Improved 3D UNet model on Rangpur GPU cluster + +import os +import torch +from torch import nn, optim +from torch.utils.data import DataLoader, random_split +from dataset import HipMRIDataset +from modules import ImprovedUNet3D +from torch.cuda.amp import autocast, GradScaler +import torch.nn.functional as F + + + +# ---------------------------- +# CONFIGURATION +# ---------------------------- +MRI_DIR = "/home/groups/comp3710/HipMRI_Study_open/semantic_MRs" +LABEL_DIR = "/home/groups/comp3710/HipMRI_Study_open/semantic_labels_only" + +EPOCHS = 15 # Increase if GPU allows +BATCH_SIZE = 1 +LR = 0.001 +DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu') + +print(f"Using device: {DEVICE}") +if DEVICE.type == "cuda": + print("CUDA available — training on GPU") +else: + print("CUDA not available — training on CPU") + +# ---------------------------- +# DATASET & DATALOADER +# ---------------------------- +dataset = HipMRIDataset(MRI_DIR, LABEL_DIR, transform=None, crop=True) +if len(dataset) == 0: + raise RuntimeError(f"No .nii.gz files found in {MRI_DIR}. Please check dataset path.") + +# 80% train, 20% val +train_size = int(0.8 * len(dataset)) +val_size = len(dataset) - train_size +train_set, val_set = random_split(dataset, [train_size, val_size]) + +train_loader = DataLoader(train_set, batch_size=BATCH_SIZE, shuffle=True, num_workers=4, pin_memory=True) +val_loader = DataLoader(val_set, batch_size=1, shuffle=False, num_workers=2, pin_memory=True) + + +# ---------------------------- +# MODEL, LOSS, OPTIMIZER +# ---------------------------- +model = ImprovedUNet3D().to(DEVICE) + +# ← BETTER: Combined loss function +class DiceCELoss(nn.Module): + def __init__(self): + super().__init__() + self.ce = nn.CrossEntropyLoss() + + def forward(self, pred, target): + ce_loss = self.ce(pred, target) + + # Dice loss + pred_soft = torch.softmax(pred, dim=1) + dice_loss = 0 + for c in range(pred.shape[1]): + pred_c = pred_soft[:, c] + target_c = (target == c).float() + intersection = (pred_c * target_c).sum() + union = pred_c.sum() + target_c.sum() + dice_loss += 1 - (2 * intersection + 1e-6) / (union + 1e-6) + dice_loss /= pred.shape[1] + + return ce_loss + dice_loss + +criterion = DiceCELoss() +optimizer = optim.Adam(model.parameters(), lr=LR, weight_decay=1e-4) +scheduler = optim.lr_scheduler.ReduceLROnPlateau( + optimizer, mode='max', factor=0.5, patience=3 +) +def multiclass_dice(pred, target, num_classes=6, eps=1e-6): + """Compute mean Dice coefficient across all classes.""" + dice_scores = [] + for c in range(num_classes): + pred_c = (pred == c).float() + target_c = (target == c).float() + intersection = (pred_c * target_c).sum() + union = pred_c.sum() + target_c.sum() + dice = (2 * intersection + eps) / (union + eps) + dice_scores.append(dice) + return torch.mean(torch.stack(dice_scores)) + + +# ---------------------------- +# TRAINING LOOP +# ---------------------------- +scaler = GradScaler(enabled=(DEVICE.type == "cuda")) +for epoch in range(EPOCHS): + model.train() + epoch_loss = 0.0 + epoch_dice = 0.0 + + for batch_idx, (img, label) in enumerate(train_loader): + img = img.to(DEVICE) + label = label.to(DEVICE).long() # important for CrossEntropyLoss + + optimizer.zero_grad() + with autocast(enabled=(DEVICE.type == "cuda")): + output = model(img) + if output.shape[2:] != label.shape[1:]: + label = F.interpolate( + label.unsqueeze(1).float(), + size=output.shape[2:], + mode="nearest" + ).squeeze(1).long() + loss = criterion(output, label) + + scaler.scale(loss).backward() + scaler.step(optimizer) + scaler.update() + torch.cuda.empty_cache() + + + with torch.no_grad(): + preds = torch.argmax(output, dim=1) + dice = multiclass_dice(preds, label) + + epoch_loss += loss.item() + epoch_dice += dice.item() + + print(f"Epoch [{epoch+1}/{EPOCHS}] Batch [{batch_idx+1}/{len(train_loader)}] " + f"Loss: {loss.item():.4f} | Dice: {dice.item():.4f}") + + # summary for the epoch + avg_loss = epoch_loss / len(train_loader) + avg_dice = epoch_dice / len(train_loader) + print(f"\n Epoch [{epoch+1}/{EPOCHS}] Train Loss: {avg_loss:.4f} | Train Dice: {avg_dice:.4f}") + + # ---- VALIDATION ---- + model.eval() + val_dice = 0.0 + with torch.no_grad(): + for img, label in val_loader: + img = img.to(DEVICE) + label = label.to(DEVICE).long() + output = model(img) + + if output.shape[2:] != label.shape[1:]: + label = F.interpolate( + label.unsqueeze(1).float(), + size=output.shape[2:], + mode="nearest" + ).squeeze(1).long() + + preds = torch.argmax(output, dim=1) + val_dice += multiclass_dice(preds, label).item() + val_dice_avg = val_dice / len(val_loader) + scheduler.step(val_dice_avg) + for g in optimizer.param_groups: + g['lr'] = max(g['lr'] * 0.8, 1e-5) + + # Print final Dice after last epoch + if epoch == EPOCHS - 1: + print(f"\n Final Training Dice Coefficient: {avg_dice:.4f}") + + + +# ---------------------------- +# SAVE MODEL +# ---------------------------- +os.makedirs("models", exist_ok=True) +torch.save(model.state_dict(), "models/improved_unet3d.pth") +print("Training complete! Model saved to models/improved_unet3d.pth") diff --git a/recognition/README.md b/recognition/README.md new file mode 100644 index 000000000..272bba4fa --- /dev/null +++ b/recognition/README.md @@ -0,0 +1,20 @@ +# Pattern Analysis +Pattern Analysis of various datasets by COMP3710 students in 2025 at the University of Queensland. + +We create pattern recognition and image processing library for Tensorflow (TF), PyTorch or JAX. + +This library is created and maintained by The University of Queensland [COMP3710](https://my.uq.edu.au/programs-courses/course.html?course_code=comp3710) students. + +The library includes the following implemented in Tensorflow: +* fractals +* recognition problems + +In the recognition folder, you will find many recognition problems solved including: +* segmentation +* classification +* graph neural networks +* StyleGAN +* Stable diffusion +* transformers +etc. +