src/channels.py

"""
The single source of truth for what every input channel means.

There are three different input encodings in this release. Getting them mixed up
is the easiest way to get silently wrong dose predictions, so they are all
spelled out here and every checkpoint records which one it expects
(``ckpt["channel_spec"]``).

    SOURCE_8CH    the 8-channel encoding the *pretrained* checkpoints expect
    HAN_5CH       the 5-channel encoding the head-and-neck models expect
    PANCREAS_5CH  the 5-channel encoding the pancreas model expects

The first five channels line up positionally across all three, which is why
transfer learning works at all. They are NOT numerically interchangeable --
see the per-spec notes below.

Dose-valued channels are always divided by ``dose_div_factor`` (10), so a 70 Gy
prescription appears as 7.0.
"""

# --------------------------------------------------------------------------
# 8-channel source-domain encoding (the *_pretrained.pt checkpoints)
# --------------------------------------------------------------------------
# Reproduced from the pretraining data loader (GDP-HMM / RT dose challenge
# codebase), whose pipeline follows Gao et al., "Flexible-CM GAN", CVPR 2023.
SOURCE_8CH = [
    ("comb_optptv",   "Optimisation PTVs (the plan's OPTName structures), each "
                      "multiplied by its prescription/10, merged by voxel-wise max."),
    ("comb_ptv",      "Planning PTVs (the StructName structures), each multiplied "
                      "by its prescription/10, merged by voxel-wise max."),
    ("comb_oar",      "All OARs merged into one label-encoded channel: each mask "
                      "scaled by 1 + 4*i/30 (i = index in the site's OAR list), "
                      "merged by max."),
    ("body",          "Binary external / BODY mask."),
    ("img",           "CT, clipped to [-1000, 1000] HU and divided by 500 -> [-2, 2]."),
    ("beam_plate",    "Beam geometry map: for each gantry angle, rays are traced "
                      "from the source through the PTV surface, dilated/eroded and "
                      "weighted by inverse-square distance; summed over all angles."),
    ("angle_plate",   "2D gantry-angle map replicated over the +/-5 slices around "
                      "the isocentre in Z; zero on every other slice."),
    ("prompt_extend", "Plan metadata [isVMAT, n_PTVs, site, cohort] broadcast to a "
                      "full volume as four equal blocks stacked along Z."),
]

# --------------------------------------------------------------------------
# 5-channel head-and-neck encoding (c3d/mednext/swinunetr _finetuned/_fromscratch)
# --------------------------------------------------------------------------
# Channels 5-7 of the source encoding are DROPPED, not zero-filled: the released
# HaN models physically have a 5-channel input stem. The private HaN dataset has
# no beam geometry or plan-prompt metadata, so those three channels carried no
# signal and were removed rather than fed as zeros.
HAN_5CH = [
    ("comb_optptv",   "PTVs multiplied by prescription/10, merged by max. "
                      "NOTE: byte-identical to channel 1 -- the HaN loader builds "
                      "the optimisation and planning PTV sets from the same "
                      "StructName masks. Kept as a separate channel only to match "
                      "the pretrained model's layout."),
    ("comb_ptv",      "Same as channel 0. See the note above."),
    ("comb_oar",      "23 HaN OARs merged into one channel, each mask scaled by "
                      "1 + 4*i/30 for i = 1..23, merged by max -> values in "
                      "[1.13, 4.07]."),
    ("body",          "Binary external / BODY mask."),
    ("img",           "CT, clipped to [-1000, 1000] HU and divided by 500 -> [-2, 2]."),
]

# --------------------------------------------------------------------------
# 5-channel pancreas encoding (pancreas_c3d_finetuned / _fromscratch)
# --------------------------------------------------------------------------
# Same five slots as HAN_5CH, but channels 0, 1 and 2 are computed differently.
# A HaN-encoded volume fed to the pancreas model (or vice versa) will run
# without error and give wrong answers.
PANCREAS_5CH = [
    ("comb_optptv",   "PTVs with every overlapping OAR subtracted "
                      "(ptv * (1 - any_oar)), then multiplied by prescription/10. "
                      "Models the sub-volume that can realistically reach the "
                      "prescription. NOT equal to channel 1."),
    ("comb_ptv",      "PTVs multiplied by min(prescription/10, 3.0) -- the value is "
                      "CAPPED at 30 Gy. A 40 Gy PTV appears as 3.0, not 4.0."),
    ("comb_oar",      "14 PancVMAT OARs merged into one channel with two tiers: "
                      "dose-limiting OARs (Stomach, Bowel_Small, Bowel_Large and "
                      "their PRVs) scale to 2.0 + 0.5*i/14 -> [2.0, 2.5]; all other "
                      "OARs scale to 0.5 + 0.3*i/14 -> [0.5, 0.8]."),
    ("body",          "Binary external / BODY mask."),
    ("img",           "CT, clipped to [-1000, 1000] HU and divided by 500 -> [-2, 2]."),
]

CHANNEL_SPECS = {
    "source_8ch": SOURCE_8CH,
    "han_5ch": HAN_5CH,
    "pancreas_5ch": PANCREAS_5CH,
}


def names(spec):
    """Channel names for a spec key, in order."""
    return [n for n, _ in CHANNEL_SPECS[spec]]


def describe(spec):
    """Human-readable description of a channel spec."""
    lines = [f"{spec} ({len(CHANNEL_SPECS[spec])} channels)"]
    for i, (name, doc) in enumerate(CHANNEL_SPECS[spec]):
        lines.append(f"  ch{i}  {name}")
        lines.append(f"        {doc}")
    return "\n".join(lines)


def pad_to_source_8ch(x):
    """
    Zero-pad a 5-channel batch [B, 5, Z, Y, X] to the 8-channel source layout.

    This lets a *_pretrained.pt checkpoint run on data prepared by the released
    5-channel loaders, but the result is not a meaningful dose prediction: the
    pretrained model was trained with real beam_plate / angle_plate / prompt
    channels and relies on them. Use it for shape debugging and for warm-starting
    your own fine-tuning, not for evaluation.
    """
    import torch

    if x.shape[1] != 5:
        raise ValueError(f"expected 5 input channels, got {x.shape[1]}")
    pad = torch.zeros(x.shape[0], 3, *x.shape[2:], dtype=x.dtype, device=x.device)
    return torch.cat((x, pad), dim=1)


if __name__ == "__main__":
    for key in CHANNEL_SPECS:
        print(describe(key))
        print()