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):
mermaid
1flowchart TD
2 A[DICOM series x N]--> B["2.5D clip sampler (prev/center/next)"]3 B --> C["ConvNeXtV2-Tiny backbone (shared, runs once for the whole batch)"]4 C --> D["Projection: Linear -> LayerNorm -> GELU -> Dropout (embed_dim=384)"]
5 D --> E["Label-Aware Series Attention (12 learned per-label queries)"]6 F["Metadata: Anatomical_Plane / Fluid_Sensitive / Fat_Suppression"]--> E
7 G["Soft priors: PLANE_PRIOR / FLUID_PRIOR / FAT_PRIOR (additive bias, not a hard filter)"]--> E
8 E --> H["Label Token Transformer (2 layers, 8 heads)"]9 H --> I["Vectorized Label Heads (12 parallel MLPs)"]10 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:
Read the pixel array, apply RescaleSlope/RescaleIntercept.
Invert if PhotometricInterpretation == MONOCHROME1.
Percentile windowing (default 1st–99th percentile) → normalize to
[0, 1].
Resize to img_size (default 224×224).
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.
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
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-infersAnatomical_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.
