Model Architecture

Model Architecture

The primary model is a 3D Attention U-Net with squeeze-and-excitation channel attention and spatial attention gates. A lighter 2D U-Net is also included for quick experiments on synthetic data.

---

3D Attention U-Net (src/segmentation/unet3d.py)

from src.segmentation import AttentionUNet3D

model = AttentionUNet3D(
    in_channels=1,
    num_classes=4,
    base_filters=32,
    depth=4,
    dropout=0.1
)
# Input:  (B, 1, D, H, W)
# Output: (B, 4, D, H, W)  — 4 classes: background, glioma, meningioma, pituitary

Encoder

Each encoder level consists of:

1. Conv3DBlock — two 3×3×3 convolutions with BatchNorm + ReLU
2. Squeeze-and-Excitation (SE) block — channel attention: global average pool → FC → sigmoid → scale
3. Down3D — MaxPool3d(kernel_size=(1,2,2)) to downsample H and W only, preserving depth

The (1,2,2) kernel is intentional. The model processes pseudo-3D inputs with small depth (D=4), so standard (2,2,2) pooling would collapse the depth dimension by the third encoder level.

Bottleneck

Conv3DBlock at the deepest level with higher dropout.

Decoder

Each decoder level:

1. Up3D — ConvTranspose3d(kernel_size=(1,2,2)) to upsample H and W only
2. Spatial attention gate on the skip connection — suppresses irrelevant spatial locations before concatenation
3. Conv3DBlock — fuse upsampled features + attended skip connection

Output

1×1×1 convolution → 4-channel logits → argmax at inference time.

---

Pseudo-3D Training

The Kaggle dataset contains 2D MRI slices, not full 3D volumes. To use the 3D model, each slice is stacked D_FRAMES times to create a (B, 1, D, H, W) volume:

D_FRAMES = 4  # configurable in scripts/train3d.py
images_3d = images.unsqueeze(2).repeat(1, 1, D_FRAMES, 1, 1)

This is a pragmatic bridge — the model learns spatial context within a slice at multiple scales. With true volumetric data (e.g., BraTS), D_FRAMES would be replaced by real depth.

---

2D U-Net (src/segmentation/unet.py)

from src.segmentation import UNet

model = UNet(in_channels=1, n_classes=1, base_filters=32, depth=4)
# Input:  (B, 1, H, W)
# Output: (B, 1, H, W)  — binary logits

Used for quick training on synthetic data (scripts/train.py). Binary output — tumor vs background.

---

Parameter Count

Model	Depth	Base Filters	Parameters
2D U-Net	4	32	~7M
3D Attention U-Net	3	32	~15M
3D Attention U-Net	4	32	~31M

---

Loss Functions (src/segmentation/losses_advanced.py)

Loss	Purpose
WeightedDiceLoss	Per-class Dice with class weights — handles imbalance
FocalLoss	Down-weights easy background pixels
BoundaryLoss	Penalizes imprecise tumor edges
HybridLoss	α·WeightedDice + β·Focal + γ·Boundary

Default weights in train3d.py: α=0.5, β=0.3, γ=0.2. Class weights for 4-class: [0.01, 0.4, 0.3, 0.29] (background gets near-zero weight).