classify_tissue.py

#!/usr/bin/env python3
"""Map Allen ontology structure ids (voxel_count.csv) -> 9 tissue classes (map.csv).

Deterministic, hierarchy-driven. Each structure inherits the tissue value of the
*nearest* anchor on its structure_id_path (deepest anchor wins). Anchors are the
ontology's own gray/white/ventricle/surface/vessel division nodes, keyed by
acronym so the script is robust to id changes.

Writes LUTs (space-separated `id value`):
  tissue_map.lut            - all nodes
  tissue_map_annotated.lut  - only annotated=True (voxel-bearing) nodes
  tissue_map_background.lut  - visualization aid: background-class nodes -> 1, else 0
"""
import csv
import sys

VOXEL_CSV = "voxel_count.csv"

# tissue values (map.csv column 1)
#   0 background  1 cerebrospinalFluid  2 grayMatter  3 whiteMatter
#   4 ventricle   5 cerebellarGrayMatter 6 cerebellarWhiteMatter 7 brainStem
#   8 deepGrayMatter
ANCHORS = {
    # forebrain (cerebrum): real gray/white/ventricle split
    "FGM": 2,  # gray matter of forebrain (default: cortex; deep nuclei override below)
    "FWM": 3,  # white matter of forebrain
    "FV": 4,   # ventricles of forebrain
    # deep gray matter (subcortical/diencephalic nuclei) -> overrides FGM(2). These
    # are siblings of cortex (Cx) under Tel/Die, so they're deeper than FGM.
    # Cortex stays 2; hippocampus is filed under FSS (surface) so it is untouched.
    "CN": 8,     # cerebral nuclei: basal ganglia (caudate/putamen/GPe/GPi/NAC),
                 #   amygdala, claustrum, basal forebrain, septum, BNST
    "THM": 8,    # thalamus (+ habenula, pineal, zona incerta)
    "SubTH": 8,  # subthalamus (subthalamic nucleus)
    "HTH": 8,    # hypothalamus (preoptic/supraoptic/tuberal/mammillary)
    # forebrain surface structures = macroscopic landmarks. Default gray (gyri),
    # but sulci/fissures are CSF-filled subarachnoid space, peri-ventricular too.
    "FSS": 2,    # surface structures of forebrain (gyri/lobules + surface nuclei)
    "CeS": 1,    # cerebral sulci & fissures -> CSF (overrides FSS)
    "ASFV": 1,   # adjoining structures of forebrain ventricles -> CSF
    "fbv": 0,  # blood vessels of forebrain
    "FTS": 0,  # transient (developmental) structures of forebrain
    # midbrain -> brainStem (map.csv has no separate midbrain GM/WM)
    "MGM": 7, "MWM": 7, "MV": 4, "MSS": 7, "mbv": 0,
    # hindbrain -> brainStem by default; cerebellum (deeper) overrides to 5
    "HGM": 7, "HV": 4, "HSS": 7, "hbv": 0,
    # white matter of hindbrain -> cerebellarWM: its bulk voxels are the
    # cerebellar internal white matter (arbor vitae). CB is filed under the gray
    # branch (HGM), so the cerebellar internal WM has no dedicated node and lands
    # in this catch-all. Spatial check on annotation_full.nii.gz: 100% of HWM
    # voxels lie inside the cerebellum bbox; 85.5% are closer to the cerebellum
    # centroid than the brainstem. (icp/mcp peduncles under here inherit 6.)
    "HWM": 6,
    "HTS": 0,  # transient structures of hindbrain
    # cerebellum: cortex + deep nuclei (under CB) and surface lobules (under
    # CbSS) are all gray; cerebellar fissures are CSF.
    "CB": 5,
    "CbSS": 5,   # surface structures of cerebellum (lobes/lobules) -> overrides HSS
    "cbf": 1,    # cerebellar fissures -> CSF (overrides CbSS)
    # cerebellar peduncles filed under MWM (midbrain) -> cerebellarWM; deeper than
    # MWM so deepest-wins overrides brainStem(7). icp/mcp are under HWM and
    # inherit HWM's 6, so they need no explicit anchor.
    "scp": 6,    # superior cerebellar peduncle (under MWM)
    "xscp": 6,   # decussation of scp (under MWM; 0 voxels, anatomically cerebellar)
    # spinal cord: not a brain tissue class
    "SpC": 0,
}
DEFAULT = 0  # bare container nodes (NP, NT, Br, F, M, H) and anything unanchored


def load_rows(path):
    with open(path, newline="") as f:
        return list(csv.DictReader(f))


def main():
    rows = load_rows(VOXEL_CSV)

    # acronym -> id (warn on collision)
    acc = {}
    for r in rows:
        a = r["acronym"]
        if a in acc and acc[a] != r["id"]:
            print(f"WARN: duplicate acronym {a} ids {acc[a]} / {r['id']}", file=sys.stderr)
        acc.setdefault(a, r["id"])

    # resolve anchor acronyms -> anchor ids
    anchor_ids = {}
    for a, v in ANCHORS.items():
        if a in acc:
            anchor_ids[acc[a]] = v
        else:
            print(f"WARN: anchor acronym {a} not found in ontology", file=sys.stderr)

    # classify: nearest anchor on path (leaf -> root) wins
    def classify(row):
        path = [p for p in row["structure_id_path"].split("/") if p]
        for pid in reversed(path):
            if pid in anchor_ids:
                return anchor_ids[pid]
        return DEFAULT

    mapping = {}        # id -> value
    annotated_ids = []
    for r in rows:
        v = classify(r)
        mapping[r["id"]] = v
        if r["annotated"] == "True":
            annotated_ids.append(r["id"])

    # write LUTs, numeric-sorted by id
    def write_lut(path, ids):
        with open(path, "w") as f:
            for i in sorted(ids, key=int):
                f.write(f"{i} {mapping[i]}\n")

    write_lut("tissue_map.lut", list(mapping))
    write_lut("tissue_map_annotated.lut", annotated_ids)

    # secondary visualization map: structures that classified to background(0)
    # -> 1, everything else -> 0. Lets the otherwise-invisible background-class
    # structures be highlighted when applied to a label volume.
    with open("tissue_map_background.lut", "w") as f:
        for i in sorted(mapping, key=int):
            f.write(f"{i} {1 if mapping[i] == 0 else 0}\n")

    # audit
    names = {0: "background", 1: "cerebrospinalFluid", 2: "grayMatter",
             3: "whiteMatter", 4: "ventricle", 5: "cerebellarGrayMatter",
             6: "cerebellarWhiteMatter", 7: "brainStem", 8: "deepGrayMatter"}
    print(f"total nodes: {len(mapping)}   annotated: {len(annotated_ids)}")
    print("\nper-class counts (all nodes / annotated):")
    ann_set = set(annotated_ids)
    for v in range(9):
        allc = sum(1 for x in mapping.values() if x == v)
        annc = sum(1 for i in annotated_ids if mapping[i] == v)
        print(f"  {v} {names[v]:<22} {allc:>5} / {annc:>3}")

    # flag annotated structures that fell to background (likely misclassified)
    by_id = {r["id"]: r for r in rows}
    bg = [i for i in annotated_ids if mapping[i] == 0]
    if bg:
        print(f"\nWARN: {len(bg)} annotated (voxel-bearing) structures -> background(0):")
        for i in bg:
            r = by_id[i]
            print(f"  {i} {r['acronym']:<8} {r['name']}")
    else:
        print("\nOK: every annotated structure mapped to a tissue class 1-8.")

    if not any(v == 6 for v in mapping.values()):
        print("NOTE: class 6 (cerebellarWhiteMatter) has no members "
              "(ontology splits cerebellum into cortex + deep nuclei, both gray).")


if __name__ == "__main__":
    main()