make_coarse.py

#!/usr/bin/env python3
"""Derive coarser tissue schemes from the 9-class classification + priors.

Collapses the existing 8 tissue classes (in classify.nii.gz / tissue_map.lut)
into smaller label sets, then regenerates the LUTs, the relabeled volume,
and the probabilistic priors for each — reusing the exact same smoothing
(2 mm FWHM, DiscreteGaussian) and per-voxel normalization as make_priors.py.

Schemes:

  4class   CSF / GM / WM / DeepGM            (brainStem -> WM, ventricle -> CSF)
  3class   CSF / GM / WM                     (brainStem -> WM, deepGM -> GM)
  cbmerge  CSF / GM / WM / ventricle /       (full 9-class granularity, but
           cerebellum / brainStem / deepGM    cerebellar GM+WM merged into one)

For each scheme <tag> it writes:
  map_<tag>.csv
  tissue_map_<tag>.lut            (all nodes)
  tissue_map_<tag>_annotated.lut  (voxel-bearing nodes)
  classify_<tag>.nii.gz           (uint8)
  prior_<tag>_0N.nii.gz           (float32, 0..1, N = 1..K)
"""
import csv
import math

import numpy as np
import SimpleITK as sitk

SRC = "classify.nii.gz"          # 9-class hard segmentation (incl. mask CSF fill)
LUT = "tissue_map.lut"
LUT_ANN = "tissue_map_annotated.lut"
FWHM = 2.0                       # mm (matches make_priors.py)
EPS = 1e-4

# source classes: 0 bg 1 CSF 2 GM 3 WM 4 ventricle 5 cerebGM 6 cerebWM 7 brainStem 8 deepGM
SCHEMES = {
    "4class": {
        "collapse": {0: 0, 1: 1, 2: 2, 3: 3, 4: 1, 5: 2, 6: 3, 7: 3, 8: 4},
        "classes": {
            0: ("background", "unlabelled voxels"),
            1: ("cerebrospinalFluid", "CSF incl. ventricles"),
            2: ("grayMatter", "cortical + cerebellar gray matter"),
            3: ("whiteMatter", "cerebral + cerebellar white matter + brain stem"),
            4: ("deepGrayMatter", "deep gray matter (subcortical/diencephalic nuclei)"),
        },
    },
    "3class": {
        "collapse": {0: 0, 1: 1, 2: 2, 3: 3, 4: 1, 5: 2, 6: 3, 7: 3, 8: 2},
        "classes": {
            0: ("background", "unlabelled voxels"),
            1: ("cerebrospinalFluid", "CSF incl. ventricles"),
            2: ("grayMatter", "all gray matter (cortex, cerebellar, deep nuclei)"),
            3: ("whiteMatter", "all white matter + brain stem"),
        },
    },
    # full 9-class granularity, but cerebellar GM (5) + WM (6) merged into one
    # cerebellum label; remaining classes kept and renumbered contiguous.
    "cbmerge": {
        "collapse": {0: 0, 1: 1, 2: 2, 3: 3, 4: 4, 5: 5, 6: 5, 7: 6, 8: 7},
        "classes": {
            0: ("background", "unlabelled voxels"),
            1: ("cerebrospinalFluid", "CSF excluding the ventricles"),
            2: ("grayMatter", "cerebral gray matter (cortex)"),
            3: ("whiteMatter", "cerebral white matter"),
            4: ("ventricle", "ventricles (CSF inside the brain)"),
            5: ("cerebellum", "cerebellar gray + white matter merged"),
            6: ("brainStem", "brain stem"),
            7: ("deepGrayMatter", "deep gray matter (subcortical/diencephalic nuclei)"),
        },
    },
}


def write_map_csv(tag, classes):
    with open(f"map_{tag}.csv", "w", newline="") as f:
        w = csv.writer(f)
        w.writerow(["value", "shortName", "description"])
        for v in sorted(classes):
            w.writerow([v, classes[v][0], classes[v][1]])


def compose_lut(src_path, dst_path, collapse):
    with open(src_path) as fi, open(dst_path, "w") as fo:
        for line in fi:
            i, old = line.split()
            fo.write(f"{i} {collapse[int(old)]}\n")


def build_priors(tag, lab, img, K, names):
    sigma = FWHM / (2.0 * math.sqrt(2.0 * math.log(2.0)))
    smoothed = {}
    for c in range(1, K + 1):
        m = (lab == c).astype(np.float32)
        mi = sitk.GetImageFromArray(m)
        mi.CopyInformation(img)
        si = sitk.DiscreteGaussian(mi, variance=sigma ** 2, useImageSpacing=True)
        smoothed[c] = np.clip(sitk.GetArrayFromImage(si), 0.0, None)

    total = np.zeros_like(lab, dtype=np.float32)
    for c in range(1, K + 1):
        total += smoothed[c]
    support = total > EPS

    print(f"  {'file':<24}{'tissue':<22}{'peak':>7}{'nonzero':>12}")
    for c in range(1, K + 1):
        p = np.zeros_like(total)
        p[support] = smoothed[c][support] / total[support]
        p = p.astype(np.float32)
        out = sitk.GetImageFromArray(p)
        out.CopyInformation(img)
        path = f"prior_{tag}_{c:02d}.nii.gz"
        sitk.WriteImage(out, path)
        print(f"    {path:<22}{names[c]:<22}{p.max():>7.3f}{int((p > 0).sum()):>12}")

    # verify: in-support sum == 1, outside == 0
    s = np.zeros_like(total)
    for c in range(1, K + 1):
        s += sitk.GetArrayFromImage(sitk.ReadImage(f"prior_{tag}_{c:02d}.nii.gz"))
    ins = s[support]
    ok = abs(ins.min() - 1) < 1e-4 and abs(ins.max() - 1) < 1e-4 and s[~support].max() == 0.0
    print(f"  priors sum-in-support [{ins.min():.6f},{ins.max():.6f}] "
          f"-> {'CONSISTENT' if ok else 'INVESTIGATE'}")


def main():
    img = sitk.ReadImage(SRC)
    arr = sitk.GetArrayFromImage(img)

    for tag, spec in SCHEMES.items():
        collapse = spec["collapse"]
        classes = spec["classes"]
        K = max(classes)
        names = {v: classes[v][0] for v in classes}
        print(f"\n=== {tag}: {' / '.join(names[v] for v in range(1, K + 1))} ===")

        write_map_csv(tag, classes)
        compose_lut(LUT, f"tissue_map_{tag}.lut", collapse)
        compose_lut(LUT_ANN, f"tissue_map_{tag}_annotated.lut", collapse)

        # remap the hard volume via a small lookup (source values 0..8)
        vals = np.array([collapse[v] for v in range(arr.max() + 1)], dtype=np.uint8)
        out = vals[arr]
        oimg = sitk.GetImageFromArray(out)
        oimg.CopyInformation(img)
        sitk.WriteImage(oimg, f"classify_{tag}.nii.gz")
        uniq = np.unique(out).tolist()
        assert set(uniq) <= set(range(K + 1)), f"out-of-range labels: {uniq}"
        print(f"  classify_{tag}.nii.gz unique: {uniq}")

        build_priors(tag, out, img, K, names)


if __name__ == "__main__":
    main()