apply_classification.py

#!/usr/bin/env python3
"""Relabel annotation_full.nii.gz (Allen structure ids) -> classify.nii.gz
(tissue values 0-8) using the tissue_map.lut produced by classify_tissue.py.

The LUT is the single source of truth. Ids reach ~2.7e8 but the volume only
contains a small set of them, so we remap via np.unique + searchsorted rather
than a dense lookup array.
"""
import csv

import numpy as np
import SimpleITK as sitk

LUT = "tissue_map.lut"
SRC = "annotation_full.nii.gz"
DST = "classify.nii.gz"
VOXEL_CSV = "voxel_count.csv"
# brain mask (same grid as SRC) used to source CSF: the Allen annotation carries
# no CSF voxels, so non-brain/subarachnoid CSF (label 1) is defined spatially as
# the in-mask voxels the annotation left unlabelled.
MASK = "mni_icbm152_t1_tal_nlin_sym_09b_hires_synthstrip_mask.nii.gz"

NAMES = {0: "background", 1: "cerebrospinalFluid", 2: "grayMatter",
         3: "whiteMatter", 4: "ventricle", 5: "cerebellarGrayMatter",
         6: "cerebellarWhiteMatter", 7: "brainStem", 8: "deepGrayMatter"}


def load_lut(path):
    m = {}
    with open(path) as f:
        for line in f:
            i, v = line.split()
            m[int(i)] = int(v)
    return m


def main():
    mapping = load_lut(LUT)

    img = sitk.ReadImage(SRC)
    arr = sitk.GetArrayFromImage(img)
    print(f"read {SRC}: shape={arr.shape} dtype={arr.dtype} "
          f"size={img.GetSize()} spacing={img.GetSpacing()}")

    # sparse vectorized remap
    u = np.unique(arr)
    missing = [int(x) for x in u if int(x) not in mapping and int(x) != 0]
    if missing:
        print(f"WARN: {len(missing)} ids in volume absent from LUT -> background: "
              f"{missing[:10]}{'...' if len(missing) > 10 else ''}")
    vals = np.array([mapping.get(int(x), 0) for x in u], dtype=np.uint8)
    out = vals[np.searchsorted(u, arr)]

    # define non-brain CSF (label 1): in-mask voxels left unlabelled by the
    # annotation (sulcal / subarachnoid CSF) become CSF. Brain-mask grid must
    # match the annotation grid.
    mimg = sitk.ReadImage(MASK)
    assert mimg.GetSize() == img.GetSize(), \
        f"mask grid {mimg.GetSize()} != annotation grid {img.GetSize()}"
    m = sitk.GetArrayFromImage(mimg)
    fill = (m == 1) & (out == 0)
    out[fill] = 1
    print(f"CSF fill: {int(fill.sum())} in-mask background voxels -> CSF(1)")

    outimg = sitk.GetImageFromArray(out)
    outimg.CopyInformation(img)  # preserve spacing/origin/direction
    sitk.WriteImage(outimg, DST)
    print(f"wrote {DST}: dtype=uint8")

    # ---- verification ----
    uniq = np.unique(out)
    assert set(uniq.tolist()) <= set(range(9)), f"out has values outside 0-8: {uniq}"
    print(f"\nclassify.nii.gz unique values: {uniq.tolist()}")

    counts = np.bincount(out.ravel(), minlength=9)
    # expected per-class voxel totals by aggregating voxel_count.csv
    expected = {v: 0 for v in range(9)}
    rows = list(csv.DictReader(open(VOXEL_CSV)))
    for r in rows:
        vc = r["voxel_count"]
        if vc:
            expected[mapping.get(int(r["id"]), 0)] += int(vc)

    # The annotation volume is at 2x the count-scale of voxel_count.csv (verified
    # per-structure), so each non-empty class should share one constant ratio.
    # A constant ratio across classes is the real check: it proves the mapping is
    # internally consistent (a relabel bug would skew classes differently).
    print(f"\n{'class':<24}{'volume voxels':>14}{'csv voxels':>13}{'ratio':>8}")
    ratios = []
    for v in range(9):
        vol = int(counts[v])
        exp = expected[v]
        r = vol / exp if exp else None
        if r is not None and v != 0:
            ratios.append(r)
        print(f"  {v} {NAMES[v]:<20}{vol:>14}{exp:>13}{(f'{r:.3f}' if r else '-'):>8}")
    if ratios:
        spread = max(ratios) - min(ratios)
        # leaf structures are exactly 2.000x; small dips below 2.0 are partial-volume
        # from the half-scale csv counts (thin CSF/ventricle structures lose most).
        ok = max(ratios) <= 2.01 and min(ratios) >= 1.9
        print(f"\nper-class volume/csv ratio: {min(ratios):.3f}-{max(ratios):.3f} "
              f"(all ~2.0; <2.0 = half-scale partial-volume, not a mapping error) "
              f"-> {'CONSISTENT' if ok else 'INVESTIGATE'}")


if __name__ == "__main__":
    main()