cxas/label_mapper.py

import json
import os
from pathlib import Path
import colorcet as cc
import pprint

# 0. Load PAX Ray plus plus label dictionary
this_directory = Path(__file__).parent
paxray_labels = json.load(
    open(os.path.join(str(this_directory), "data/paxray_labels.json"), encoding="utf-8")
)
id2label_dict = paxray_labels["label_dict"]

# Atomic map for internal use
_name2id_atomic = {v: int(k) for k, v in id2label_dict.items()}

def ids(names):
    return [_name2id_atomic[n] for n in names if n in _name2id_atomic]

# 1. Atomic groups from your snippet
thoracic_spine = [
    "thoracic spine",
    *[f"vertebrae T{i}" for i in range(1, 13)],
]
all_vertebrae = [
    *[f"vertebrae L{i}" for i in range(1, 6)],
    *[f"vertebrae T{i}" for i in range(1, 13)],
    *[f"vertebrae C{i}" for i in range(1, 8)],
]
cervical_spine = [f"vertebrae C{i}" for i in range(1, 8)]
lumbar_spine   = ["lumbar spine"] + [f"vertebrae L{i}" for i in range(1, 6)]

clavicle_set = ["clavicles", "clavicle left", "clavicle right"]
scapula_set  = ["scapulas", "scapula left", "scapula right"]

rib = [
    "posterior 12th rib right", "posterior 12th rib left",
    "anterior 11th rib right",  "posterior 11th rib right",
    "anterior 11th rib left",   "posterior 11th rib left",
    "anterior 10th rib right",  "posterior 10th rib right",
    "anterior 10th rib left",   "posterior 10th rib left",
    "anterior 9th rib right",   "posterior 9th rib right",
    "anterior 9th rib left",    "posterior 9th rib left",
    "anterior 8th rib right",   "posterior 8th rib right",
    "anterior 8th rib left",    "posterior 8th rib left",
    "anterior 7th rib right",   "posterior 7th rib right",
    "anterior 7th rib left",    "posterior 7th rib left",
    "anterior 6th rib right",   "posterior 6th rib right",
    "anterior 6th rib left",    "posterior 6th rib left",
    "anterior 5th rib right",   "posterior 5th rib right",
    "anterior 5th rib left",    "posterior 5th rib left",
    "anterior 4th rib right",   "posterior 4th rib right",
    "anterior 4th rib left",    "posterior 4th rib left",
    "anterior 3rd rib right",   "posterior 3rd rib right",
    "anterior 3rd rib left",    "posterior 3rd rib left",
    "anterior 2nd rib right",   "posterior 2nd rib right",
    "anterior 2nd rib left",    "posterior 2nd rib left",
    "anterior 1st rib right",   "posterior 1st rib right",
    "anterior 1st rib left",    "posterior 1st rib left",
]
ribsuper = [
    "12th rib", "12th rib",
    "anterior 11th rib", "posterior 11th rib",
    "anterior 10th rib", "posterior 10th rib",
    "anterior 9th rib",  "posterior 9th rib",
    "anterior 8th rib",  "posterior 8th rib",
    "anterior 7th rib",  "posterior 7th rib",
    "anterior 6th rib",  "posterior 6th rib",
    "anterior 5th rib",  "posterior 5th rib",
    "anterior 4th rib",  "posterior 4th rib",
    "anterior 3rd rib",  "posterior 3rd rib",
    "anterior 2nd rib",  "posterior 2nd rib",
    "anterior 1st rib",  "posterior 1st rib",
]

diaphragm = ["diaphragm", "left hemidiaphragm", "right hemidiaphragm"]

mediastinum = [
    "cardiomediastinum",
    "upper mediastinum", "lower mediastinum",
    "anterior mediastinum", "middle mediastinum", "posterior mediastinum",
]

abdomen = [
    "stomach", "small bowel", "duodenum",
    "liver", "pancreas",
    "kidney left", "kidney right",
]

heart = [
    "heart", "heart atrium left", "heart atrium right",
    "heart myocardium", "heart ventricle left", "heart ventricle right",
]

breast = ["breast left", "breast right"]

trachea = ["trachea", "tracheal bifurcation"]

zones = [
    "right upper zone lung", "right mid zone lung", "right lung base",
    "right apical zone lung",
    "left upper zone lung",  "left mid zone lung",  "left lung base",
    "left apical zone lung",
]

lung_halves = ["right lung", "left lung"]

vessels = [
    "heart", "ascending aorta", "descending aorta", "aortic arch",
    "pulmonary artery", "inferior vena cava",
]

lobes = [
    "lung lower lobe left", "lung upper lobe left",
    "lung lower lobe right", "lung middle lobe right", "lung upper lobe right",
]

# Sided rib lists
right_ribs_anterior = [f"anterior {i}th rib right" for i in [11,10,9,8,7,6,5,4,3,2,1]]
right_ribs_posterior = ["posterior 12th rib right"] + [f"posterior {i}th rib right" for i in [11,10,9,8,7,6,5,4,3,2,1]]
left_ribs_anterior  = [f"anterior {i}th rib left"  for i in [11,10,9,8,7,6,5,4,3,2,1]]
left_ribs_posterior = ["posterior 12th rib left"]  + [f"posterior {i}th rib left"  for i in [11,10,9,8,7,6,5,4,3,2,1]]

# 1b. Pleura approximations built from existing labels
pleura_visceral = [
    "right lung", "left lung",
    "right apical zone lung", "left apical zone lung",
    "right upper zone lung", "left upper zone lung",
    "right mid zone lung", "left mid zone lung",
    "right lung base", "left lung base",
    "lung upper lobe right", "lung middle lobe right", "lung lower lobe right",
    "lung upper lobe left", "lung lower lobe left"
]

pleura_parietal_costal = rib + ["sternum", "clavicles", "scapulas"]

pleura_parietal_diaphragmatic = ["left hemidiaphragm", "right hemidiaphragm"]

pleura_parietal_mediastinal = [
    "cardiomediastinum", "heart", "aorta", "ascending aorta", "descending aorta", "aortic arch",
    "pulmonary artery", "inferior vena cava", "esophagus"
]

pleura_parietal_cervical = [
    "right apical zone lung", "left apical zone lung",
    "clavicles", "clavicle left", "clavicle right",
    "anterior 1st rib right", "anterior 1st rib left",
    "posterior 1st rib right", "posterior 1st rib left"
]

pleural_costophrenic_recess = [
    "right lung base", "left lung base",
    "left hemidiaphragm", "right hemidiaphragm",
    "posterior 12th rib right", "posterior 12th rib left",
    "anterior 11th rib right", "anterior 11th rib left"
]

pleural_cardiophrenic_recess = [
    "left hemidiaphragm", "right hemidiaphragm",
    "cardiomediastinum", "heart"
]

# 1c. Hierarchical tree with abstract parents plus pleura
KG_TREE = {
    "anatomy": [
        "bones",
        "lungs and airways",
        "mediastinum and heart",
        "diaphragm",
        "abdomen",
        "breast",
        "pleura"
    ],

    "bones": [
        "spine",
        "ribs",
        "clavicles",
        "scapulas",
        "sternum"
    ],

    "spine": ["cervical spine", "thoracic spine", "lumbar spine"],
    "cervical spine": [f"vertebrae C{i}" for i in range(1, 8)],
    "thoracic spine": [f"vertebrae T{i}" for i in range(1, 13)],
    "lumbar spine": [f"vertebrae L{i}" for i in range(1, 6)],

    "ribs": ["right ribs", "left ribs", "ribs super", "first rib"],
    "first rib": ["anterior 1st rib", "posterior 1st rib"],

    "right ribs": ["right ribs anterior", "right ribs posterior"],
    "right ribs anterior": right_ribs_anterior,
    "right ribs posterior": right_ribs_posterior,

    "left ribs": ["left ribs anterior", "left ribs posterior"],
    "left ribs anterior": left_ribs_anterior,
    "left ribs posterior": left_ribs_posterior,

    "clavicles": ["clavicle left", "clavicle right"],
    "scapulas": ["scapula left", "scapula right"],

    "lungs and airways": ["lung", "trachea"],
    "trachea": ["tracheal bifurcation"],

    "lung": [
        "right lung",
        "left lung",
        "apical zone lung",
        "upper zone lung",
        "mid zone lung",
        "lung base",
        "lung lobes by side"
    ],

    "right lung": [
        "right apical zone lung", "right upper zone lung", "right mid zone lung", "right lung base",
        "lung upper lobe right", "lung middle lobe right", "lung lower lobe right"
    ],
    "left lung": [
        "left apical zone lung", "left upper zone lung", "left mid zone lung", "left lung base",
        "lung upper lobe left", "lung lower lobe left"
    ],
    "lung lobes by side": [
        "lung upper lobe right", "lung middle lobe right", "lung lower lobe right",
        "lung upper lobe left", "lung lower lobe left"
    ],

    "mediastinum and heart": ["cardiomediastinum", "heart", "aorta", "pulmonary artery", "inferior vena cava", "esophagus"],

    "diaphragm": ["left hemidiaphragm", "right hemidiaphragm"],

    "abdomen": ["stomach", "bowel", "liver", "pancreas", "kidneys"],
    "bowel": ["small bowel", "duodenum"],
    "kidneys": ["kidney left", "kidney right"],

    "breast": ["breast left", "breast right"],
    "sternum": [],

    "pleura": ["visceral pleura approx", "parietal pleura approx", "pleural recesses approx"],
    "parietal pleura approx": ["costal pleura approx", "diaphragmatic pleura approx", "mediastinal pleura approx", "cervical pleura approx"],
    "pleural recesses approx": ["costophrenic recess approx", "cardiophrenic recess approx"]
}

# 2. label_mapper with flat id lists
label_mapper = {id2label_dict[k]: [int(k)] for k in id2label_dict.keys()}

# Original composites as flat ids
label_mapper.update({
    "thoracic spine":   ids(thoracic_spine),
    "all vertebrae":    ids(all_vertebrae),
    "cervical spine":   ids(cervical_spine),
    "lumbar spine":     ids(lumbar_spine),
    "clavicle set":     ids(clavicle_set),
    "scapula set":      ids(scapula_set),
    "ribs":             ids(rib),
    "ribs super":       ids(ribsuper),
    "diaphragm":        ids(diaphragm),
    "mediastinum":      ids(mediastinum),
    "abdomen":          ids(abdomen),
    "heart region":     ids(heart),
    "breast tissue":    ids(breast),
    "trachea":          ids(trachea),
    "lung zones":       ids(zones),
    "lung halves":      ids(lung_halves),
    "vessels":          ids(vessels),
    "lung lobes":       ids(lobes),
    "lung lobes by side": ids(["lung upper lobe right", "lung middle lobe right", "lung lower lobe right",
                               "lung upper lobe left", "lung lower lobe left"]),
})

# New abstract composites including pleura
label_mapper.update({
    "anatomy": ids(list(id2label_dict.values())),
    "bones": ids(["spine", "clavicles", "scapulas", "sternum"] + rib + ribsuper),
    "right ribs": ids(right_ribs_anterior + right_ribs_posterior),
    "left ribs": ids(left_ribs_anterior + left_ribs_posterior),
    "right ribs anterior": ids(right_ribs_anterior),
    "right ribs posterior": ids(right_ribs_posterior),
    "left ribs anterior": ids(left_ribs_anterior),
    "left ribs posterior": ids(left_ribs_posterior),
    "first rib": ids(["anterior 1st rib", "posterior 1st rib"]),
    "lungs and airways": ids(["lung"] + trachea + zones + lobes + lung_halves),
    "mediastinum and heart": ids(["cardiomediastinum", "heart", "aorta", "pulmonary artery", "inferior vena cava", "esophagus"]),
    "kidneys": ids(["kidney left", "kidney right"]),
    "bowel": ids(["small bowel", "duodenum"]),
    "visceral pleura approx": ids(pleura_visceral),
    "parietal pleura approx": ids(pleura_parietal_costal + pleura_parietal_diaphragmatic + pleura_parietal_mediastinal + pleura_parietal_cervical),
    "costal pleura approx": ids(pleura_parietal_costal),
    "diaphragmatic pleura approx": ids(pleura_parietal_diaphragmatic),
    "mediastinal pleura approx": ids(pleura_parietal_mediastinal),
    "cervical pleura approx": ids(pleura_parietal_cervical),
    "pleural recesses approx": ids(pleural_costophrenic_recess + pleural_cardiophrenic_recess),
    "costophrenic recess approx": ids(pleural_costophrenic_recess),
    "cardiophrenic recess approx": ids(pleural_cardiophrenic_recess),
})

# 4. build a unified name2id that returns flat id lists for every node
name2id = {name: [nid] for name, nid in _name2id_atomic.items()}
for k, v in label_mapper.items():
    name2id[k] = list(sorted(set(v)))

# 3. sections to flat ids
def ids_any(names):
    out = []
    for n in names:
        if n in name2id:
            out.extend(name2id[n])          # works for atomic and composite
        elif n in _name2id_atomic:
            out.append(_name2id_atomic[n])  # fallback atomic
    # unique and stable order
    return list(dict.fromkeys(out))

SECTION_TO_NAMES = {
    "abdominal": [
        "stomach", "small bowel", "duodenum",
        "liver", "pancreas", "kidney left", "kidney right"
    ],
    "cardiovascular": [
        "heart", "heart atrium left", "heart atrium right",
        "heart ventricle left", "heart ventricle right", "heart myocardium",
        "aorta", "ascending aorta", "descending aorta", "aortic arch",
        "pulmonary artery", "inferior vena cava"
    ],
    "hila and mediastinum": [
        "cardiomediastinum",
        "upper mediastinum", "lower mediastinum",
        "anterior mediastinum", "middle mediastinum", "posterior mediastinum",
        "esophagus"
    ],
    "lungs and airways": [
        "lung halves", "lung zones", "lung lobes", "trachea"
    ],
    "musculoskeletal and chest wall": [
        "spine", "ribs", "clavicles", "scapulas", "sternum", "breast tissue"
    ],
    "pleura": [
        "visceral pleura approx",
        "parietal pleura approx",
        "pleural recesses approx",
        "diaphragm"
    ],
    "tubes, catheters, and support devices": ["anatomy"],
    "other": ["anatomy"]
}

# Flatten each section through name2id into flat id lists
for section_key, names in SECTION_TO_NAMES.items():
    label_mapper[section_key] = ids_any(names)

for section_key in SECTION_TO_NAMES.keys():
    name2id[section_key] = list(dict.fromkeys(label_mapper[section_key]))

from functools import lru_cache

@lru_cache(maxsize=None)
def atomic_descendants(node):
    if node in _name2id_atomic:
        return [node]
    out = []
    if node in KG_TREE:
        for ch in KG_TREE[node]:
            out.extend(atomic_descendants(ch))
    return list(dict.fromkeys(out))

for node in KG_TREE.keys():
    atoms = atomic_descendants(node)
    if atoms:
        name2id[node] = ids(atoms)

# 5. optional edges
KG_EDGES = []
for parent, children in KG_TREE.items():
    for child in children:
        KG_EDGES.append((parent, child, "part_of"))

# 6. colors preserved
colors        = [cc.cm.glasbey_bw_minc_20(i) for i in range(len(id2label_dict))]
colors_alpha  = [[c[0], c[1], c[2], c[3] / 2] for c in colors]
category_colors = {tuple(colors[i][:3]): i for i in range(len(colors))}
category_ids    = {id2label_dict[k]: int(k) for k in id2label_dict}

# if __name__ == "__main__":
#     print("Visceral pleura ids size:", len(name2id["visceral pleura approx"]))
#     print("Parietal pleura ids size:", len(name2id["parietal pleura approx"]))
#     print("Costal pleura ids sample:", name2id["costal pleura approx"][:8])
#     print("Costophrenic recess ids sample:", name2id["costophrenic recess approx"][:8])