HMCLIP / scripts /plot_anat_cosne.py
SteveZeyuZhang's picture
Add source files and streaming backup restore instructions
b91f38a
Raw History Blame Contribute Delete
23.6 kB
#!/usr/bin/env python3
"""
Disease-colored CO-SNE for RadIR/RadIR hyperbolic embeddings.
Uses the same hyperbolic CO-SNE pipeline as scripts/plot_cosne.py, then colors
points by disease while keeping modality markers:
○ Image, △ Text
Color = disease / class
Labeling strategies (--color_by):
gt ground-truth exclusive / priority labels (see --label_mode)
inferred disease-seeded DBSCAN on high-dim hyp image feats ("反推法")
cosne2d same seeded DBSCAN but on 2D CO-SNE image coords
cluster_csv load precomputed cluster_disease from infer_disease_clusters.py
Usage (from RadIR/RadIR):
# Ground-truth exclusive (clean)
python scripts/plot_anat_cosne.py \\
--load_emb outputs/cosne_emb.npz \\
--color_by gt --label_mode exclusive --hide_other \\
--out outputs/cosne_disease_clean.png
# Reverse-inference clusters on CO-SNE disk
python scripts/plot_anat_cosne.py \\
--load_emb outputs/cosne_emb.npz \\
--load_feats outputs/tsne_feats.npz \\
--color_by inferred --hide_other \\
--out outputs/cosne_disease_inferred.png
"""
from __future__ import annotations
import argparse
import os
import sys
from collections import Counter
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import torch
ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
SCRIPTS = os.path.dirname(os.path.abspath(__file__))
RAD_IR_ROOT = os.path.join(ROOT, "Rad_IR")
for p in (ROOT, RAD_IR_ROOT, SCRIPTS):
if p not in sys.path:
sys.path.insert(0, p)
import infer_disease_clusters as idc # noqa: E402
import plot_cosne as cosne # noqa: E402
DISEASE_NAMES = idc.DISEASE_NAMES
CLASS_ORDER = ["Normal"] + DISEASE_NAMES + ["Other"]
CLASS_COLORS = {
"Normal": "#00C853",
"Pneumothorax": "#AA00FF",
"Edema": "#00B8D4",
"Pneumonia": "#D50000",
"Effusion": "#2962FF",
"Cardiomegaly": "#FF6D00",
"Atelectasis": "#FFD600",
"Other": "#BDBDBD",
}
MODALITY_MARKER = {"Image": "o", "Text": "^"}
def style_for(cls: str, modality: str) -> tuple[str, str]:
return CLASS_COLORS.get(cls, "#333333"), MODALITY_MARKER[modality]
def load_label_table(labels_csv: str) -> pd.DataFrame:
df = pd.read_csv(labels_csv)
needed = ["dicom_id", "normal_only"] + [f"{d}_pos" for d in DISEASE_NAMES]
missing = [c for c in needed if c not in df.columns]
if missing:
raise KeyError(f"labels CSV missing columns: {missing}")
return df.drop_duplicates(subset=["dicom_id"], keep="first").reset_index(drop=True)
def build_labels(
df: pd.DataFrame,
sample_idx: np.ndarray,
label_mode: str,
) -> list[str]:
"""
Map global row indices → class name.
exclusive: single-label only; else "Other"
priority: Normal if normal_only, else first disease hit, else "Other"
"""
labels: list[str] = []
for i in sample_idx:
row = df.iloc[int(i)]
hits: list[str] = []
if int(row["normal_only"]) == 1:
hits.append("Normal")
for d in DISEASE_NAMES:
if int(row[f"{d}_pos"]) == 1:
hits.append(d)
if label_mode == "exclusive":
labels.append(hits[0] if len(hits) == 1 else "Other")
elif label_mode == "priority":
if int(row["normal_only"]) == 1 and len(hits) == 1:
labels.append("Normal")
else:
assigned = None
for d in DISEASE_NAMES:
if int(row[f"{d}_pos"]) == 1:
assigned = d
break
if assigned is None and int(row["normal_only"]) == 1:
assigned = "Normal"
labels.append(assigned if assigned is not None else "Other")
else:
raise ValueError(label_mode)
return labels
def _to_plot_label(name: str) -> str:
if name in (None, "", "Unassigned", "Other"):
return "Other"
return str(name)
def labels_from_cluster_csv(
cluster_csv: str, labels_csv: str, sample_idx: np.ndarray
) -> list[str]:
"""Map sample_idx → cluster_disease via dicom_id alignment."""
gt = pd.read_csv(labels_csv).drop_duplicates(subset=["dicom_id"], keep="first")
cl = pd.read_csv(cluster_csv).drop_duplicates(subset=["dicom_id"], keep="first")
cl = cl.set_index("dicom_id")["cluster_disease"]
out: list[str] = []
for i in sample_idx:
did = str(gt.iloc[int(i)]["dicom_id"])
if did in cl.index:
out.append(_to_plot_label(cl.loc[did]))
else:
out.append("Other")
return out
def labels_inferred_highdim(
feats_npz: str,
labels_csv: str,
sample_idx: np.ndarray,
seed_mode: str,
eps: float | None,
min_samples: int,
eps_quantile: float,
min_cluster_size: int,
) -> tuple[list[str], list[dict]]:
"""
Disease-seeded DBSCAN on full-set high-dim hyp image feats, then slice
assignments to the CO-SNE subsample.
"""
data = np.load(feats_npz, allow_pickle=True)
feats = idc.prepare_image_feats(data, "hyperbolic")
df = pd.read_csv(labels_csv).drop_duplicates(subset=["dicom_id"], keep="first")
if len(df) != feats.shape[0]:
raise RuntimeError(
f"feats N={feats.shape[0]} vs labels rows={len(df)} "
"(expect same valid.jsonl order)"
)
normal_only = df["normal_only"].astype(int).to_numpy().astype(bool)
disease_mat = np.stack(
[df[f"{d}_pos"].astype(int).to_numpy().astype(bool) for d in DISEASE_NAMES],
axis=1,
)
print(f"===== Infer clusters on high-dim hyp feats ({seed_mode}) =====")
assigned, _uid, stats = idc.infer_disease_clusters(
feats,
normal_only,
disease_mat,
seed_mode=seed_mode,
eps=eps,
min_samples=min_samples,
eps_quantile=eps_quantile,
min_cluster_size=min_cluster_size,
)
labels = [_to_plot_label(assigned[int(i)]) for i in sample_idx]
return labels, stats
def labels_inferred_cosne2d(
emb_img: np.ndarray,
df: pd.DataFrame,
sample_idx: np.ndarray,
seed_mode: str,
eps: float | None,
min_samples: int,
eps_quantile: float,
min_cluster_size: int,
) -> tuple[list[str], list[dict]]:
"""
Same disease-seeded DBSCAN, but distance is Euclidean in the 2D CO-SNE
Poincaré plane (answers: does the disk layout form disease neighborhoods?).
"""
n = emb_img.shape[0]
normal_only = np.zeros(n, dtype=bool)
disease_mat = np.zeros((n, len(DISEASE_NAMES)), dtype=bool)
for k, i in enumerate(sample_idx):
row = df.iloc[int(i)]
normal_only[k] = bool(int(row["normal_only"]))
for j, d in enumerate(DISEASE_NAMES):
disease_mat[k, j] = bool(int(row[f"{d}_pos"]))
# Use Euclidean metric on 2D coords via L2-normalize trick disabled:
# cluster_subset uses cosine on L2-normalized vectors; for 2D Poincaré
# points we pass raw coords and temporarily use a Euclidean DBSCAN path.
print(f"===== Infer clusters on CO-SNE 2D ({seed_mode}) =====")
assigned, _uid, stats = infer_clusters_euclidean_2d(
emb_img,
normal_only,
disease_mat,
seed_mode=seed_mode,
eps=eps,
min_samples=min_samples,
eps_quantile=eps_quantile,
min_cluster_size=min_cluster_size,
)
labels = [_to_plot_label(a) for a in assigned]
return labels, stats
def infer_clusters_euclidean_2d(
coords: np.ndarray,
normal_only: np.ndarray,
disease_mat: np.ndarray,
seed_mode: str,
eps: float | None,
min_samples: int,
eps_quantile: float,
min_cluster_size: int,
) -> tuple[np.ndarray, np.ndarray, list[dict]]:
"""DBSCAN with Euclidean metric on 2D coords (within disease seeds)."""
from collections import defaultdict
from sklearn.cluster import DBSCAN
from sklearn.neighbors import NearestNeighbors
n = coords.shape[0]
excl = idc.exclusive_labels(normal_only, disease_mat)
candidates: list[dict] = []
class_names = ["Normal"] + DISEASE_NAMES
for name in class_names:
if seed_mode == "exclusive":
member_idx = np.where(excl == name)[0]
else:
if name == "Normal":
member_idx = np.where(normal_only)[0]
else:
j = DISEASE_NAMES.index(name)
member_idx = np.where(disease_mat[:, j])[0]
print(f" Seed {name:14s} ({seed_mode}): {member_idx.size} images", end="")
if member_idx.size < min_cluster_size:
print(" → 0 clusters")
continue
sub = coords[member_idx].astype(np.float64)
eps_used = eps
if eps_used is None:
k = min(min_samples, sub.shape[0] - 1)
if k < 1:
print(" → 0 clusters")
continue
nn = NearestNeighbors(n_neighbors=k + 1, metric="euclidean").fit(sub)
dists, _ = nn.kneighbors(sub)
eps_used = float(np.quantile(dists[:, -1], eps_quantile))
labels_db = DBSCAN(
eps=eps_used, min_samples=min_samples, metric="euclidean"
).fit_predict(sub)
n_kept = 0
for cid in sorted(set(labels_db.tolist()) - {-1}):
local = np.where(labels_db == cid)[0]
if local.size < min_cluster_size:
continue
global_idx = member_idx[local]
candidates.append(
{
"disease": name,
"global_indices": global_idx.astype(int),
"size": int(local.size),
"eps_used": float(eps_used),
}
)
n_kept += 1
print(f" → {n_kept} clusters ≥{min_cluster_size} (eps={eps_used:.4f})")
claims: dict[int, list[tuple]] = defaultdict(list)
for cid, c in enumerate(candidates):
idxs = c["global_indices"]
sub = coords[idxs]
centroid = sub.mean(axis=0)
dist = np.linalg.norm(sub - centroid, axis=1)
for k, i in enumerate(idxs):
is_excl = 1 if excl[i] == c["disease"] else 0
claims[int(i)].append(
(is_excl, c["size"], -float(dist[k]), cid, c["disease"])
)
assigned = np.array(["Unassigned"] * n, dtype=object)
cluster_uid = np.full(n, -1, dtype=int)
for i, opts in claims.items():
opts.sort(reverse=True)
_, _, _, cid, disease = opts[0]
assigned[i] = disease
cluster_uid[i] = cid
stats = []
for cid, c in enumerate(candidates):
idxs = c["global_indices"]
kept = idxs[assigned[idxs] == c["disease"]]
stats.append(
{
"cluster_uid": cid,
"disease": c["disease"],
"size_raw": c["size"],
"size_kept": int(kept.size),
"eps_used": c["eps_used"],
}
)
return assigned, cluster_uid, stats
def print_cluster_stats(stats: list[dict], labels: list[str]) -> None:
print("----- kept clusters -----")
by_d: dict[str, int] = Counter()
for s in stats:
if s.get("size_kept", 0) > 0:
print(
f" uid={s['cluster_uid']:3d} {s['disease']:14s} "
f"kept={s['size_kept']:4d} raw={s['size_raw']:4d} "
f"eps={s['eps_used']:.4f}"
)
by_d[s["disease"]] += s["size_kept"]
print("----- subsample assignment -----")
counts = Counter(labels)
for cls in CLASS_ORDER:
if counts.get(cls, 0):
print(f" {cls:14s} {counts[cls]}")
labeled = sum(v for k, v in counts.items() if k != "Other")
print(f" labeled={labeled}/{len(labels)} Other={counts.get('Other', 0)}")
def run_or_load_cosne(args) -> tuple[np.ndarray, np.ndarray, np.ndarray, float]:
"""
Returns emb_img [M,2], emb_txt [M,2], sample_idx [M], curv.
"""
if args.load_emb:
print(f"Loading CO-SNE coords: {args.load_emb}")
data = np.load(args.load_emb, allow_pickle=True)
emb_img = np.asarray(data["emb_img"])
emb_txt = np.asarray(data["emb_txt"])
sample_idx = np.asarray(data["sample_idx"], dtype=np.int64)
curv = float(data["curv"]) if "curv" in data.files else 1.0
if emb_img.shape[0] != emb_txt.shape[0] or emb_img.shape[0] != len(sample_idx):
raise ValueError("emb_img / emb_txt / sample_idx length mismatch")
return emb_img, emb_txt, sample_idx, curv
if not args.load_feats:
raise ValueError("Provide --load_emb or --load_feats")
print(f"Loading hyperbolic feats: {args.load_feats}")
feats = np.load(args.load_feats, allow_pickle=True)
if "hyp_img" not in feats.files:
raise KeyError("Need hyp_img / hyp_txt in feats npz")
hyp_img = torch.from_numpy(feats["hyp_img"])
hyp_txt = torch.from_numpy(feats["hyp_txt"])
curv = float(feats["curv"]) if "curv" in feats.files else 1.0
max_samples = None if args.max_samples <= 0 else args.max_samples
hyp_img, hyp_txt, sample_idx = cosne.subsample_pair(
hyp_img, hyp_txt, max_samples, args.seed
)
print(f" subsampled to {len(sample_idx)} pairs, curv={curv:.4f}")
with torch.no_grad():
poi_img = cosne.lorentz_space_to_poincare(hyp_img.float(), curv).numpy()
poi_txt = cosne.lorentz_space_to_poincare(hyp_txt.float(), curv).numpy()
X = np.concatenate([poi_img, poi_txt], axis=0)
model = cosne.COSNE(
n_components=2,
perplexity=args.perplexity,
n_iter=args.n_iter,
learning_rate=args.learning_rate,
learning_rate_h=args.learning_rate_h,
early_exaggeration=args.early_exaggeration,
student_t_gamma=args.student_t_gamma,
seed=args.seed,
device=args.cosne_device,
verbose=True,
)
Y = model.fit_transform(X)
n = poi_img.shape[0]
emb_img, emb_txt = Y[:n], Y[n:]
if args.save_emb:
os.makedirs(
os.path.dirname(os.path.abspath(args.save_emb)) or ".", exist_ok=True
)
np.savez(
args.save_emb,
emb_img=emb_img,
emb_txt=emb_txt,
sample_idx=sample_idx,
curv=curv,
)
print(f"Saved CO-SNE coords to {args.save_emb}")
return emb_img, emb_txt, sample_idx, curv
def plot_disease_poincare(
emb_img: np.ndarray,
emb_txt: np.ndarray,
labels: list[str],
out_path: str,
title: str,
dpi: int,
hide_other: bool,
max_pair_lines: int,
) -> None:
n = emb_img.shape[0]
assert len(labels) == n
fig, ax = plt.subplots(figsize=(9, 9))
ax.add_patch(plt.Circle((0, 0), 1.0, color="black", fill=False, linewidth=1.2))
for r in (0.25, 0.5, 0.75):
ax.add_patch(
plt.Circle((0, 0), r, color="gray", fill=False, linewidth=0.4, alpha=0.35)
)
# optional pair connectors (only for non-Other if hide_other)
n_lines = min(max_pair_lines, n)
drawn = 0
for i in range(n):
if drawn >= n_lines:
break
if hide_other and labels[i] == "Other":
continue
ax.plot(
[emb_img[i, 0], emb_txt[i, 0]],
[emb_img[i, 1], emb_txt[i, 1]],
color="gray",
alpha=0.12,
linewidth=0.5,
zorder=0,
)
drawn += 1
# draw Other first (background), then each disease
draw_order = ["Other"] + [c for c in CLASS_ORDER if c != "Other"]
for cls in draw_order:
if hide_other and cls == "Other":
continue
mask = np.array([l == cls for l in labels], dtype=bool)
if not mask.any():
continue
for modality, emb in (("Image", emb_img), ("Text", emb_txt)):
color, marker = style_for(cls, modality)
is_other = cls == "Other"
ax.scatter(
emb[mask, 0],
emb[mask, 1],
c=color,
marker=marker,
s=18 if is_other else (32 if modality == "Image" else 42),
alpha=0.35 if is_other else 0.88,
edgecolors="none" if is_other else "black",
linewidths=0.0 if is_other else 0.3,
label=f"{cls}-{modality} (n={int(mask.sum())})",
zorder=1 if is_other else (2 if modality == "Image" else 3),
)
# radius summary on labeled (non-Other) points
labeled_mask = np.array([l != "Other" for l in labels], dtype=bool)
if labeled_mask.any():
r_img = np.linalg.norm(emb_img[labeled_mask], axis=1).mean()
r_txt = np.linalg.norm(emb_txt[labeled_mask], axis=1).mean()
subtitle = f"labeled mean radius Image={r_img:.3f} Text={r_txt:.3f}"
else:
subtitle = "no exclusive labels in subsample"
ax.set_title(f"{title}\n{subtitle}")
ax.set_xlim(-1.05, 1.05)
ax.set_ylim(-1.05, 1.05)
ax.set_aspect("equal")
ax.legend(loc="upper left", bbox_to_anchor=(1.02, 1.0), fontsize=8, frameon=True)
ax.axis("off")
fig.tight_layout()
os.makedirs(os.path.dirname(os.path.abspath(out_path)) or ".", exist_ok=True)
fig.savefig(out_path, dpi=dpi, bbox_inches="tight")
if out_path.lower().endswith(".png"):
fig.savefig(out_path[:-4] + ".pdf", dpi=dpi, bbox_inches="tight")
plt.close(fig)
print(f"Saved disease CO-SNE to {out_path}")
def parse_args():
p = argparse.ArgumentParser(
description="Disease-colored CO-SNE (hyperbolic) for RadIR"
)
p.add_argument(
"--load_emb",
default=None,
help="Cached CO-SNE coords from plot_cosne.py (--save_emb)",
)
p.add_argument(
"--load_feats",
default=None,
help="Hyperbolic feats npz (required for --color_by inferred)",
)
p.add_argument(
"--labels_csv",
default=os.path.join(ROOT, "outputs/tsne_disease_labels.csv"),
)
p.add_argument(
"--cluster_csv",
default=os.path.join(ROOT, "outputs/tsne_cluster_disease.csv"),
help="Precomputed cluster_disease CSV for --color_by cluster_csv",
)
p.add_argument(
"--jsonl",
default=os.path.join(ROOT, "dataset/valid.jsonl"),
help="Only used to sanity-check length vs labels (optional)",
)
p.add_argument(
"--color_by",
choices=["gt", "inferred", "cosne2d", "cluster_csv"],
default="gt",
help="gt=ground-truth; inferred=反推 on high-dim; "
"cosne2d=反推 on CO-SNE plane; cluster_csv=load CSV",
)
p.add_argument(
"--label_mode",
choices=["exclusive", "priority"],
default="exclusive",
help="For --color_by gt only",
)
p.add_argument(
"--seed_mode",
choices=["exclusive", "inclusive"],
default="exclusive",
help="Disease seed mode for inferred / cosne2d",
)
p.add_argument("--eps", type=float, default=None, help="DBSCAN eps (auto if None)")
p.add_argument("--min_samples", type=int, default=5)
p.add_argument("--eps_quantile", type=float, default=0.5)
p.add_argument("--min_cluster_size", type=int, default=8)
p.add_argument(
"--hide_other",
action="store_true",
help="Do not draw Other/Unassigned points",
)
p.add_argument("--max_samples", type=int, default=500)
p.add_argument("--perplexity", type=float, default=30.0)
p.add_argument("--n_iter", type=int, default=800)
p.add_argument("--learning_rate", type=float, default=0.05)
p.add_argument("--learning_rate_h", type=float, default=5.0)
p.add_argument("--early_exaggeration", type=float, default=4.0)
p.add_argument("--student_t_gamma", type=float, default=0.1)
p.add_argument("--seed", type=int, default=42)
p.add_argument("--cosne_device", default="cpu")
p.add_argument("--max_pair_lines", type=int, default=40)
p.add_argument(
"--out",
default=os.path.join(ROOT, "outputs/cosne_disease.png"),
)
p.add_argument("--save_emb", default=None)
p.add_argument("--dpi", type=int, default=150)
return p.parse_args()
def main():
args = parse_args()
if args.load_emb is None and args.load_feats is None:
default_emb = os.path.join(ROOT, "outputs/cosne_emb.npz")
default_feats = os.path.join(ROOT, "outputs/tsne_feats.npz")
if os.path.isfile(default_emb):
args.load_emb = default_emb
elif os.path.isfile(default_feats):
args.load_feats = default_feats
else:
raise FileNotFoundError(
"Need --load_emb or --load_feats (or outputs/cosne_emb.npz)"
)
emb_img, emb_txt, sample_idx, curv = run_or_load_cosne(args)
df = load_label_table(args.labels_csv)
if len(df) < int(sample_idx.max()) + 1:
raise RuntimeError(
f"labels CSV has {len(df)} rows but sample_idx max={sample_idx.max()}"
)
stats: list[dict] = []
if args.color_by == "gt":
labels = build_labels(df, sample_idx, args.label_mode)
print(f"===== Label mode: gt/{args.label_mode} (n={len(labels)}) =====")
for cls in CLASS_ORDER:
c = Counter(labels).get(cls, 0)
if c:
print(f" {cls:14s} {c}")
title = (
f"CO-SNE + GT ({args.label_mode}) | hyperbolic | "
f"n={len(labels)} | curv={curv:.3f}"
)
elif args.color_by == "cluster_csv":
labels = labels_from_cluster_csv(args.cluster_csv, args.labels_csv, sample_idx)
print(f"===== Color by cluster CSV: {args.cluster_csv} =====")
for cls in CLASS_ORDER:
c = Counter(labels).get(cls, 0)
if c:
print(f" {cls:14s} {c}")
title = (
f"CO-SNE + inferred clusters (CSV) | hyperbolic | "
f"n={len(labels)} | curv={curv:.3f}"
)
elif args.color_by == "inferred":
feats_path = args.load_feats or os.path.join(ROOT, "outputs/tsne_feats.npz")
if not os.path.isfile(feats_path):
raise FileNotFoundError(f"--color_by inferred needs feats: {feats_path}")
labels, stats = labels_inferred_highdim(
feats_path,
args.labels_csv,
sample_idx,
seed_mode=args.seed_mode,
eps=args.eps,
min_samples=args.min_samples,
eps_quantile=args.eps_quantile,
min_cluster_size=args.min_cluster_size,
)
print_cluster_stats(stats, labels)
title = (
f"CO-SNE + inferred (high-dim {args.seed_mode}) | hyperbolic | "
f"n={len(labels)} | curv={curv:.3f}"
)
elif args.color_by == "cosne2d":
labels, stats = labels_inferred_cosne2d(
emb_img,
df,
sample_idx,
seed_mode=args.seed_mode,
eps=args.eps,
min_samples=args.min_samples,
eps_quantile=args.eps_quantile,
min_cluster_size=args.min_cluster_size,
)
print_cluster_stats(stats, labels)
title = (
f"CO-SNE + inferred (2D disk {args.seed_mode}) | hyperbolic | "
f"n={len(labels)} | curv={curv:.3f}"
)
else:
raise ValueError(args.color_by)
plot_disease_poincare(
emb_img,
emb_txt,
labels,
args.out,
title=title,
dpi=args.dpi,
hide_other=args.hide_other,
max_pair_lines=args.max_pair_lines,
)
if __name__ == "__main__":
main()