Temporal Cubic PatchGAN
- The paper introduces TCuP-GAN, a novel GAN model that effectively captures inter-slice temporal dependencies using ConvLSTM and a 3D PatchGAN discriminator.
- It employs a Temporal Cubic U-Net generator and a multi-class segmentation approach to enforce local spatial coherence and robust volume-to-volume segmentation.
- Quantitative results on BraTS 2023 benchmarks show high Dice scores for adult glioma, while challenges remain for pediatric and Sub-Saharan Africa cohorts due to data scarcity and small lesion sizes.
Temporal Cubic PatchGAN (TCuP-GAN) is a generative adversarial network (GAN) architecture for automated three-dimensional (3D) biomedical image segmentation. It combines a generative feature learning framework with convolutional long short-term memory (ConvLSTM) modules for capturing cross-slice temporal dependencies, and it utilizes a 3D PatchGAN discriminator to enforce local spatial coherence by adversarial training. TCuP-GAN is designed for robust multi-class volume-to-volume segmentation, as demonstrated on four benchmarks from the 2023 Brain Tumor Segmentation (BraTS) Challenge, including Adult Glioma, Meningioma, Pediatric Tumors, and the Sub-Saharan Africa cohort (Mantha et al., 2023).
1. Model Architecture
1.1 Generator: Temporal Cubic U-Net
The generator, termed "Temporal Cubic U-Net," employs a U-Net style encoder–decoder structure instantiated as follows:
- Input: Image cube of dimensions (Depth × Channels × Height × Width), where the four channels comprise the T1, T2, T1ce, and FLAIR MRI sequences.
- Encoder: Five down-sampling stages. At each level :
- ConvLSTM2D (kernel ) with filters (). Each ConvLSTM2D layer accumulates spatial and inter-slice "temporal" context along the depth dimension.
- 2D convolution (, stride 1), instance normalization, leaky ReLU activation.
- 2D max pooling () reduces spatial resolution.
- Output after encoder: .
- Decoder: Five up-sampling blocks, each consisting of:
- 2D transposed convolution to double spatial resolution,
- ConvLSTM2D with filters (matching encoder depth), skip-connection from corresponding encoder activations,
- Instance normalization, leaky ReLU activation.
- Final layer: 2D convolution (, stride 1) with sigmoid activation.
- Total generator parameters: Approximately 2.3 million.
1.2 Discriminator: 3D PatchGAN
The discriminator is a 3D PatchGAN:
- Input: Concatenated MRI image cube and segmentation mask; size 0 with 1 classes for multi-class segmentation.
- Four convolutional blocks:
- Each: 3D convolution (kernel 2, stride 1), instance normalization, tanh, and dropout (0.2).
- Filter sizes: 16, 32, 64, 128; height/width are down-sampled by stride, depth is preserved.
- Final layer: Conv3D (3), sigmoid activation, outputs 4 patch authenticity map.
1.3 Role of ConvLSTM2D
ConvLSTM2D layers in both encoder and decoder enable learning of joint 2D spatial and inter-slice ("temporal") representations, providing each slice context from neighboring slices via hidden state propagation. This design is crucial for modeling the volumetric continuity intrinsic to medical image data (Mantha et al., 2023).
2. Mathematical Objectives
2.1 Adversarial Loss
The adversarial loss implements a PatchGAN objective:
5
2.2 Supervised Segmentation Loss
A scaled, class-weighted binary cross-entropy is used:
6
with 7; 8.
Additionally, a multi-class Dice loss can be used:
9
where 0, 1 are predicted and ground-truth one-hot values.
2.3 Combined Generator Objective
The generator is trained with:
2
where 3.
The discriminator minimizes:
4
3. Training and Implementation Details
3.1 Data Preprocessing and Augmentation
- Intensity normalization of each MRI sequence by 99th-percentile.
- Concatenation of four modality channels, resizing 5, 6 to 7.
- On-the-fly addition of Gaussian noise (8, 9) to non-background voxels.
- Segmentation masks are one-hot encoded into 3 classes and resized.
3.2 Training Schedule
- Optimizer: Adam (0, 1) for both generator and discriminator.
- Learning rates: 2 initial 3, decayed by 4 every 5 epochs; 5, identical decay.
- Batch size: 2 volumes.
- Trained from scratch for 30 epochs on Adult Glioma. Transfer learning to Meningioma, Pediatric, SSA cohorts for 30 epochs at constant 6.
4. Quantitative Results and Comparative Performance
Validation performance on BraTS 2023 test sets is summarized by LesionWise Dice similarity (7) and 8 Hausdorff distance (9):
| Dataset | Mean D_{lw} (WT/TC/ET) | Mean H_{95,lw} (WT [mm]/TC/ET) |
|---|---|---|
| Adult Glioma (GLI) | 0.83 / 0.76 / 0.76 | 29.3 / 43.5 / 35.6 |
| Meningioma (MEN) | 0.67 / 0.68 / 0.68 | 86.8 / 82.2 / 81.8 |
| Pediatric (PED) | 0.74 / 0.66 / 0.45 | 48.8 / 58.7 / 152.9 |
| Sub-Saharan Africa | 0.67 / 0.64 / 0.63 | 76.0 / 74.7 / 72.0 |
0: LesionWise Dice score (WT: Whole Tumor, TC: Tumor Core, ET: Enhancing Tumor).
Strong Dice (1) is attained on the GLI cohort. A reduction in ET segmentation performance is observed on smaller PED and SSA sets, reflecting data scarcity and small lesion sizes (Mantha et al., 2023).
5. Ablation Analysis, Limitations, and Failure Modes
5.1 ConvLSTM Modules
Ablation replacing ConvLSTM2D with standard 2D conv layers produces a 2 reduction in WT Dice on GLI (from 3 to 4). This demonstrates the criticality of temporal context modeling for 3D segmentation performance.
5.2 Limitations
- Cases lacking predicted TC or ET completely yield 5 for those classes.
- Small, disconnected false-positive lesions reduce mean Dice.
- Elevated 6 in PED ET class is linked to both tiny ground-truth ET volumes and large false positives—distinguishing architectural limitations from data scarcity is challenging.
- Median Dice remains 7 on most cohorts, indicating the robustness of the primary method.
6. Extensions and Future Directions
Generalization to other 3D biomedical segmentation tasks, such as multi-organelle electron microscopy, is achievable by modifying input/output channels and thresholding. Incorporating more complex objective functions—adversarial plus multi-class Dice and focal loss—facilitates stabilization in class-imbalanced scenarios. Extending the framework with attention gates or 3D ConvLSTM modules could allow learning of fully volumetric temporal features. Integration with citizen-science platforms for human-in-the-loop annotation enables rapid expert correction and downstream model fine-tuning (Mantha et al., 2023).