leminhhung0101/knee-model

RSNA Knee MRI — Model Architecture

1. Pipeline overview

DICOM series (multiple per study)
        │
        ▼
  2.5D clip (prev, center, next slice)      ← each series is cut into 3-channel clips
        │
        ▼
  ConvNeXtV2-Tiny (shared backbone)         ← ALL clips in the batch pass through ONCE
        │
        ▼
  + Metadata embedding                       ← Anatomical_Plane / Fluid_Sensitive / Fat_Suppression
        │
        ▼
  Label-Aware Multi-Series Attention         ← 12 learned queries (1 per label), attend over clips
        │
        ▼
  12 label tokens → Label Token Transformer  ← labels exchange information with each other
        │
        ▼
  Vectorized Label Heads (12 heads in parallel)
        │
        ▼
  12 probabilities (sigmoid), each in [0, 1]

Mermaid version (renders in viewers that support it):

flowchart TD
    A[DICOM series x N] --> B["2.5D clip sampler (prev/center/next)"]
    B --> C["ConvNeXtV2-Tiny backbone (shared, runs once for the whole batch)"]
    C --> D["Projection: Linear -> LayerNorm -> GELU -> Dropout (embed_dim=384)"]
    D --> E["Label-Aware Series Attention (12 learned per-label queries)"]
    F["Metadata: Anatomical_Plane / Fluid_Sensitive / Fat_Suppression"] --> E
    G["Soft priors: PLANE_PRIOR / FLUID_PRIOR / FAT_PRIOR (additive bias, not a hard filter)"] --> E
    E --> H["Label Token Transformer (2 layers, 8 heads)"]
    H --> I["Vectorized Label Heads (12 parallel MLPs)"]
    I --> J["12 sigmoid probabilities"]

---

2. Input & the 12 output labels

Each study (one patient / one scan) contains multiple series (Sagittal / Coronal / Axial, with or without fat suppression, fluid-sensitive sequences, etc.). The model predicts 12 independent probabilities:

#	Label	Meaning
1	ACL	Anterior cruciate ligament
2	MCL	Medial collateral ligament
3	Medial Meniscus	Medial meniscus tear
4	Lateral Meniscus	Lateral meniscus tear
5	Medial OA	Medial compartment osteoarthritis
6	Lateral OA	Lateral compartment osteoarthritis
7	PF OA	Patellofemoral osteoarthritis
8	Effusion	Joint effusion
9	Synovitis	Synovial inflammation
10	Baker's	Baker's (popliteal) cyst
11	Contusion	Bone contusion / bone marrow edema
12	Fracture	Fracture

Missing labels (NaN) are never treated as 0 — they are masked out of the loss everywhere (both during training and when merging manual labels).

---

3. Data preparation (before the model sees anything)

• train_series.csv is many-to-one with train.csv (1 Study → many Series), so it's split into train_series_level.csv (1 row per series) and train_study_level.csv (1 row per study, used for training).
• Manual labels (a manual-review CSV, merged by StudyInstanceUID) override the official labels wherever the value is valid (0/1); invalid or missing values become NaN and stay missing — never inferred.
• Data is split with GroupKFold on StudyInstanceUID — never split by slice or by series, to avoid leakage between train/val.
• Series metadata is aggregated up to the study level: series counts per plane, fluid-sensitive/fat-suppression counts and ratios, and counts of important plane×fluid×fat combinations (e.g. Sagittal + Fluid + Fat).

---

4. DICOM preprocessing → 2.5D clip

For each DICOM slice:

1. Read the pixel array, apply RescaleSlope/RescaleIntercept.
2. Invert if PhotometricInterpretation == MONOCHROME1.
3. Percentile windowing (default 1st–99th percentile) → normalize to [0, 1].
4. Resize to img_size (default 224×224).
5. Per-slice z-score (mean/std of that individual slice), then clip to [-5, 5].

Slices within a series are sorted by their real physical position (ImagePositionPatient projected onto the slice-plane normal), not just InstanceNumber (which is only used as a fallback when position data is missing).

A 2.5D clip = 3 consecutive slices [prev, center, next] sampled around a center slice → a [3, H, W] tensor, treated as a 3-channel image by a 2D backbone (no 3D convolutions involved).

Horizontal flip augmentation is deliberately not used, because flipping left/right would change medial/lateral semantics (medial vs. lateral meniscus, etc.) — this is an intentional design constraint, not an omission.

Series selection when a study has more series than max_series (default 8): the selection preserves diversity — the best series per plane first, then the best series per unmet (fluid, fat) combination, and the remaining slots filled by a score (log(n_slices) plus bonuses for a valid plane / fluid / fat). This is purely an input-filtering step, not part of the model's backward pass.

---

5. Backbone: ConvNeXtV2-Tiny

• Model: convnextv2_tiny.fcmae_ft_in22k_in1k (via timm), ImageNet pretrained, global_pool="avg", original classifier removed (num_classes=0).
• Every clip from every series and every study in a batch is concatenated into one large tensor and passed through the backbone exactly once (a key GPU optimization) — per-study separation only happens later, at the attention step, using offsets to slice the batch back apart.
• Gradient checkpointing is enabled to save VRAM.
• After the backbone: Linear → LayerNorm → GELU → Dropout projects down to embed_dim (default 384) — this is the feature vector for a single clip.

---

6. Label-Aware Multi-Series Attention

This is the core of the architecture — instead of average-pooling across series, each of the 12 labels gets its own learned query vector, which learns which series matter for that specific finding.

• Query: [12, embed_dim], learned (one query per label).

• Key/Value: linear projections of the clip features, after adding a metadata embedding (plane embedding + fluid/fat linear projections) — so attention is aware of which plane a clip belongs to, and whether it's fluid-sensitive/fat-suppressed.

• Additive bias on the attention scores, for every (label, clip) pair, made of two parts:

  • Domain-knowledge soft priors (fixed constants, not learned) — e.g. ACL is best seen on Sagittal, PF OA is best seen on Axial:

    Label	Sagittal	Coronal	Axial
    ACL	1.00	0.65	0.40
    MCL	0.65	1.00	0.45
    Medial/Lateral Meniscus	0.90–1.00	0.90–1.00	0.60–0.65
    Medial/Lateral OA	0.85–0.90	1.00	0.40
    PF OA	0.90	0.55	1.00
    Effusion / Synovitis	0.90	0.85	0.90
    Baker's	0.95	0.70	1.00
    Contusion / Fracture	1.00	0.90–0.95	0.80

    (there are equivalent FLUID_PRIOR and FAT_SUPPRESSION_PRIOR tables — one coefficient per label indicating how much fluid-sensitive / fat-suppressed sequences help detect it.)

  • Learned bias (plane_bias, fluid_bias, fat_bias) — lets the model deviate from the prior when the actual data suggests otherwise.

  → The priors are only an additive suggestion (soft bias), NOT a hard filter: the model can still learn to go against the prior when needed.

• Output: for each study, 12 label_feat vectors (one per label), each a weighted (attention-based) combination of every clip in that study, with weights depending on both the image content and the plane/sequence of that clip.

---

7. Label Token Transformer

• A learned per-label embedding (label_embedding) is added to label_feat.
• These 12 "label tokens" pass through a small TransformerEncoder (default 2 layers, 8 heads, pre-norm, GELU) — the purpose is to let labels exchange information with each other before independent classification (e.g. joint effusion and synovitis tend to co-occur, an ACL tear is often accompanied by bone marrow edema, etc.).

---

8. Vectorized Label Heads

• Functionally equivalent to 12 independent MLP heads (LayerNorm → Linear → GELU → Dropout → Linear), but implemented as batched 3D tensors + einsum instead of 12 separate modules in a ModuleList — faster, with less Python overhead, and mathematically identical to running 12 separate heads.
• Output: 12 logits (pre-sigmoid) → sigmoid produces the final 12 probabilities in [0, 1].

---

9. Loss function

• Robust Asymmetric Loss (ASL), masked for missing (NaN) labels:
  • Each label has its own gamma_neg (label-specific negative focusing) — rare/hard labels (e.g. Fracture, Baker's) get penalized more heavily for false negatives.
  • label_weights (computed from prevalence: 1/√prevalence, clipped to [1, 3], normalized by the mean) rebalance rare vs. common labels.
• Pairwise ranking loss: for each label, positive/negative score pairs within a batch are compared, penalizing cases where a positive score isn't higher than a negative one — this complements ASL and improves the ranking quality (useful for AUC).
• Total loss = 0.85 × ASL + 0.15 × ranking loss.

---

10. Training setup

Component	Default value
Optimizer	AdamW, backbone LR 2e-5, head LR 1e-4, weight decay 1e-4
Scheduler	Cosine with warmup (1 epoch)
Epochs	15
Batch size	4 (gradient accumulation ×2 → effective batch size 8)
Mixed precision	AMP (enabled by default)
EMA	decay 0.9997, EMA weights used for validation/inference
Grad clip	1.0
K-fold	GroupKFold on StudyInstanceUID, 5 folds
DICOM cache	mmap + uint8 (quantized from [-5,5] → [0,255]), written incrementally per series

---

11. Files in this repo

File	Role
Original training script	Prepares the data, builds the DICOM cache, trains + validates across K-folds, saves best_model.pth / last_checkpoint.pth, and produces submission.csv in Kaggle's required format.
infer_folder.py	A standalone inference script — takes any folder of DICOM series (no Kaggle-specific structure required), auto-discovers studies/series, auto-infers Anatomical_Plane/Fluid_Sensitive/Fat_Suppression from DICOM headers (can be overridden with a metadata CSV), loads a trained checkpoint, runs the exact architecture described above, and writes the 12 label probabilities to a CSV. Configured by editing plain variables in the "USER CONFIG" block at the top of the file — no CLI needed.

infer_folder.py reuses exactly the architecture described in sections 5–8 (a 1:1 copy of the LabelAwareSeriesAttention, LabelTokenTransformer, VectorizedLabelHeads, and RSNAKneeModel classes) to guarantee compatibility with trained checkpoints — it only differs in how metadata is obtained (inferred from DICOM headers instead of read from train_series.csv/test_series.csv) and drops the multi-process DICOM caching layer, which isn't needed when testing just a handful of studies.