Download scripts/plot_anat_cosne.py from AIGeeksGroup/HMCLIP: direct link, hf CLI and curl.
- Browser
- Download file 23.6 kB
-
https://huggingface.co/AIGeeksGroup/HMCLIP/resolve/main/scripts/plot_anat_cosne.py
- Command line
-
hf download hf://AIGeeksGroup/HMCLIP/scripts/plot_anat_cosne.py
-
curl -L -o plot_anat_cosne.py https://huggingface.co/AIGeeksGroup/HMCLIP/resolve/main/scripts/plot_anat_cosne.py
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() | |