| |
| |
| """ |
| Real vs. synthetic disease-mask comparison for AURAD. |
| Nature / NPG (ggsci) style. |
| |
| Two improvements over the original two-panel t-SNE: |
| |
| 1) SINGLE shared embedding, plotted in ONE panel: |
| color = disease |
| marker = source (real = filled dot, syn = open ring) |
| so you can directly see whether each disease's real & syn points |
| co-localize. |
| |
| 2) Better disease separation, via (pick with EMBED_METHOD): |
| "tsne" - PCA -> t-SNE (unsupervised) |
| "umap" - PCA -> UMAP (unsupervised, tighter clusters) |
| "umap_supervised" - PCA -> UMAP supervised on DISEASE labels only |
| (source/real-syn is NEVER given to the embedding, |
| so real<->syn overlap stays an emergent signal) |
| Plus an optional Gaussian blur on the masks (BLUR_SIGMA) so that |
| "nearby location" becomes a smooth similarity instead of 0/1 overlap. |
| |
| 3) Quantitative co-clustering evidence (COMPUTE_METRICS), per disease: |
| - C2ST AUC : classifier two-sample test, ~0.5 = real/syn indistinguishable |
| - MMD^2 (RBF) : distribution discrepancy, lower = closer |
| - norm. centroid d : Mahalanobis-normalized real<->syn centroid distance |
| - disease silhouette: how separable the diseases are (combined real+syn) |
| Saved to <prefix>_metrics.csv and printed. |
| |
| JSON record example (list of these, or one JSON per line): |
| { |
| "file_name": "chestdet/dataset/test/39796/39796.png", |
| "mask": "chestdet/dataset/test/39796/mask_Consolidation_39796.png", |
| "attn_list": [["Consolidation", "chestdet/dataset/test/39796/39796_Consolidation.png"]], |
| ... |
| } |
| |
| Real mask = REAL_ROOT / record["mask"] (REAL_MASK_SOURCE = "mask") |
| Synth mask = SYN_ROOT / attn_list[i][1] (SYN_MASK_SOURCE = "attn") |
| """ |
|
|
| import os |
| import json |
| import argparse |
|
|
| import numpy as np |
| from PIL import Image |
| from scipy.ndimage import gaussian_filter |
| from scipy.spatial.distance import cdist |
|
|
| from sklearn.decomposition import PCA |
| from sklearn.manifold import TSNE |
| from sklearn.metrics import silhouette_score |
| from sklearn.linear_model import LogisticRegression |
| from sklearn.preprocessing import StandardScaler |
| from sklearn.pipeline import make_pipeline |
| from sklearn.model_selection import cross_val_score, StratifiedKFold |
|
|
| import matplotlib as mpl |
| import matplotlib.pyplot as plt |
| from matplotlib.lines import Line2D |
|
|
| |
| JSON_PATH = "data.json" |
| REAL_ROOT = "real_masks" |
| SYN_ROOT = "syn_masks" |
|
|
| |
| |
| |
| |
| REAL_MASK_SOURCE = "mask" |
| SYN_MASK_SOURCE = "attn" |
|
|
| TARGET_DISEASES = [ |
| "Cardiomegaly", |
| "Pneumothorax", |
| "Consolidation", |
| "Effusion", |
| "Atelectasis", |
| ] |
| ATTN_KEY = "attn_list" |
|
|
| |
| MASK_SIZE = 64 |
| BLUR_SIGMA = 2.0 |
|
|
| |
| EMBED_METHOD = "umap" |
| PCA_DIM = 50 |
| TSNE_PERPLEXITY = 30 |
| UMAP_NEIGHBORS = 30 |
| UMAP_MIN_DIST = 0.40 |
| UMAP_SPREAD = 1.30 |
| RANDOM_STATE = 42 |
|
|
| |
| SYN_MARKER = "x" |
| |
| |
| DRAW_KDE = False |
| ALLOW_MULTI = False |
| MAX_PER_GROUP = None |
|
|
| |
| COMPUTE_METRICS = True |
|
|
| OUT_PREFIX = "real_vs_syn" |
| |
|
|
| |
| NPG_PALETTE = ["#E64B35", "#4DBBD5", "#00A087", "#3C5488", "#F39B7F"] |
|
|
|
|
| def set_nature_style(): |
| mpl.rcParams.update({ |
| "font.family": "sans-serif", |
| "font.sans-serif": ["Arial", "Helvetica", "DejaVu Sans"], |
| "font.size": 8, |
| "axes.linewidth": 0.8, |
| "axes.spines.top": False, |
| "axes.spines.right": False, |
| "axes.titlesize": 9, |
| "axes.titleweight": "bold", |
| "axes.labelsize": 8, |
| "xtick.major.width": 0.8, |
| "ytick.major.width": 0.8, |
| "xtick.major.size": 3, |
| "ytick.major.size": 3, |
| "legend.frameon": False, |
| "legend.fontsize": 7.5, |
| "svg.fonttype": "none", |
| "pdf.fonttype": 42, |
| }) |
|
|
|
|
| |
| def load_mask_vector(path, size, blur_sigma=0.0): |
| """Load mask -> grayscale -> resize -> (optional blur) -> /255 -> flatten.""" |
| with Image.open(path) as im: |
| im = im.convert("L").resize((size, size), Image.BILINEAR) |
| arr = np.asarray(im, dtype=np.float32) / 255.0 |
| if blur_sigma and blur_sigma > 0: |
| arr = gaussian_filter(arr, sigma=blur_sigma) |
| return arr.reshape(-1) |
|
|
|
|
| def parse_attn(attn_list): |
| if not isinstance(attn_list, (list, tuple)): |
| return |
| for entry in attn_list: |
| if isinstance(entry, (list, tuple)) and len(entry) >= 1: |
| name = entry[0] |
| path = entry[1] if len(entry) >= 2 else None |
| else: |
| name, path = entry, None |
| yield name, path |
|
|
|
|
| def record_items(rec, target_set, allow_multi, mask_source): |
| """[(relative_mask_path, disease), ...] for one record per mask_source.""" |
| diseases = [(n, p) for n, p in parse_attn(rec.get(ATTN_KEY, [])) |
| if n in target_set] |
| if not diseases: |
| return [] |
|
|
| if mask_source == "attn": |
| return [(p, n) for n, p in diseases if p] |
|
|
| rel = rec.get(mask_source) |
| if not rel: |
| return [] |
| labels = [n for n, _ in diseases] |
| labels = labels if allow_multi else labels[:1] |
| return [(rel, n) for n in labels] |
|
|
|
|
| def collect(records, root, source_tag, target_set, size, allow_multi, |
| mask_source, blur_sigma): |
| feats, diseases, missing, no_label = [], [], 0, 0 |
| for rec in records: |
| items = record_items(rec, target_set, allow_multi, mask_source) |
| if not items: |
| no_label += 1 |
| continue |
| for rel, dis in items: |
| mpath = os.path.join(root, rel) |
| if not os.path.isfile(mpath): |
| missing += 1 |
| continue |
| try: |
| vec = load_mask_vector(mpath, size, blur_sigma) |
| except Exception as e: |
| print(f" [warn] failed to read {mpath}: {e}") |
| missing += 1 |
| continue |
| feats.append(vec) |
| diseases.append(dis) |
| print(f" [{source_tag}] loaded={len(feats)} missing_files={missing} " |
| f"no_target_label={no_label}") |
| return feats, diseases |
|
|
|
|
| def subsample(feats, diseases, sources, max_per_group, seed): |
| if not max_per_group: |
| return feats, diseases, sources |
| rng = np.random.default_rng(seed) |
| keep, groups = [], {} |
| for i, (d, s) in enumerate(zip(diseases, sources)): |
| groups.setdefault((d, s), []).append(i) |
| for idxs in groups.values(): |
| if len(idxs) > max_per_group: |
| idxs = list(rng.choice(idxs, size=max_per_group, replace=False)) |
| keep.extend(idxs) |
| keep.sort() |
| return ([feats[i] for i in keep], |
| [diseases[i] for i in keep], |
| [sources[i] for i in keep]) |
|
|
|
|
| def load_records(path): |
| """Robust loader: whole-file JSON (array or dict) OR JSONL (one record/line).""" |
| with open(path, "r") as f: |
| text = f.read() |
| |
| try: |
| data = json.loads(text) |
| if isinstance(data, dict): |
| return list(data.values()) |
| if isinstance(data, list): |
| return data |
| return [data] |
| except json.JSONDecodeError: |
| pass |
| |
| records = [] |
| for line in text.splitlines(): |
| line = line.strip() |
| if line: |
| records.append(json.loads(line)) |
| return records |
|
|
|
|
| |
| def embed(X, diseases, method, seed, |
| n_neighbors=UMAP_NEIGHBORS, min_dist=UMAP_MIN_DIST, spread=UMAP_SPREAD): |
| """PCA -> {t-SNE | UMAP | supervised UMAP}. Returns (N,2) array.""" |
| n_comp = min(PCA_DIM, X.shape[0] - 1, X.shape[1]) |
| X_red = PCA(n_components=n_comp, random_state=seed).fit_transform(X) |
|
|
| if method == "tsne": |
| perp = min(TSNE_PERPLEXITY, max(5, (X.shape[0] - 1) // 3)) |
| print(f" t-SNE (perplexity={perp})") |
| emb = TSNE(n_components=2, perplexity=perp, init="pca", |
| learning_rate="auto", random_state=seed).fit_transform(X_red) |
| return emb, X_red |
|
|
| try: |
| import umap |
| except ImportError: |
| raise SystemExit("UMAP not installed. Run: pip install umap-learn\n" |
| "or set EMBED_METHOD = 'tsne'.") |
|
|
| if method == "umap": |
| print(f" UMAP (n_neighbors={n_neighbors}, min_dist={min_dist}, spread={spread})") |
| reducer = umap.UMAP(n_neighbors=n_neighbors, min_dist=min_dist, |
| spread=spread, random_state=seed) |
| emb = reducer.fit_transform(X_red) |
| elif method == "umap_supervised": |
| print(f" UMAP supervised on DISEASE only " |
| f"(n_neighbors={n_neighbors}, min_dist={min_dist}, spread={spread})") |
| y = np.array([TARGET_DISEASES.index(d) for d in diseases]) |
| reducer = umap.UMAP(n_neighbors=n_neighbors, min_dist=min_dist, |
| spread=spread, random_state=seed) |
| emb = reducer.fit_transform(X_red, y=y) |
| else: |
| raise SystemExit(f"Unknown EMBED_METHOD: {method}") |
| return emb, X_red |
|
|
|
|
| |
| def rbf_mmd2(A, B): |
| """Biased RBF-MMD^2 with the median-distance heuristic bandwidth.""" |
| if len(A) < 2 or len(B) < 2: |
| return float("nan") |
| XX = cdist(A, A, "sqeuclidean") |
| YY = cdist(B, B, "sqeuclidean") |
| XY = cdist(A, B, "sqeuclidean") |
| pooled = np.concatenate([XX[np.triu_indices_from(XX, 1)], |
| YY[np.triu_indices_from(YY, 1)], |
| XY.ravel()]) |
| med = np.median(pooled) |
| gamma = 1.0 / (med + 1e-12) |
| k = lambda D: np.exp(-gamma * D) |
| mmd2 = k(XX).mean() + k(YY).mean() - 2.0 * k(XY).mean() |
| return float(max(mmd2, 0.0)) |
|
|
|
|
| def c2st_auc(A, B, seed): |
| """Classifier two-sample test: 5-fold CV AUC for real-vs-syn. ~0.5 = good.""" |
| n = min(len(A), len(B)) |
| if n < 5: |
| return float("nan") |
| Xc = np.vstack([A, B]) |
| yc = np.concatenate([np.zeros(len(A)), np.ones(len(B))]).astype(int) |
| folds = min(5, np.bincount(yc).min()) |
| if folds < 2: |
| return float("nan") |
| clf = make_pipeline(StandardScaler(), |
| LogisticRegression(max_iter=2000)) |
| cv = StratifiedKFold(n_splits=folds, shuffle=True, random_state=seed) |
| return float(cross_val_score(clf, Xc, yc, cv=cv, scoring="roc_auc").mean()) |
|
|
|
|
| def norm_centroid_dist(A, B): |
| """||mu_A - mu_B|| / sqrt(0.5*(tr cov_A + tr cov_B)). Scale-free.""" |
| if len(A) < 2 or len(B) < 2: |
| return float("nan") |
| d = np.linalg.norm(A.mean(0) - B.mean(0)) |
| spread = np.sqrt(0.5 * (A.var(0).sum() + B.var(0).sum())) + 1e-12 |
| return float(d / spread) |
|
|
|
|
| def compute_metrics(X_red, diseases, sources, seed): |
| rows = [] |
| for d in TARGET_DISEASES: |
| m = diseases == d |
| A = X_red[m & (sources == "real")] |
| B = X_red[m & (sources == "syn")] |
| rows.append({ |
| "disease": d, |
| "n_real": int(len(A)), |
| "n_syn": int(len(B)), |
| "c2st_auc": c2st_auc(A, B, seed), |
| "mmd2": rbf_mmd2(A, B), |
| "norm_centroid_dist": norm_centroid_dist(A, B), |
| }) |
| |
| dis_int = np.array([TARGET_DISEASES.index(d) for d in diseases]) |
| sil = float(silhouette_score(X_red, dis_int)) if len(set(dis_int)) > 1 else float("nan") |
| return rows, sil |
|
|
|
|
| def print_metrics(rows, sil): |
| print("\n" + "=" * 64) |
| print(f"{'disease':14s} {'n_real':>6s} {'n_syn':>6s} " |
| f"{'C2ST':>6s} {'MMD^2':>8s} {'cdist':>7s}") |
| print("-" * 64) |
| aucs = [] |
| for r in rows: |
| auc = r["c2st_auc"] |
| if not np.isnan(auc): |
| aucs.append(auc) |
| print(f"{r['disease']:14s} {r['n_real']:6d} {r['n_syn']:6d} " |
| f"{auc:6.3f} {r['mmd2']:8.4f} {r['norm_centroid_dist']:7.3f}") |
| print("-" * 64) |
| if aucs: |
| print(f"mean C2ST AUC = {np.mean(aucs):.3f} (0.5 = real/syn indistinguishable)") |
| print(f"disease silhouette (real+syn) = {sil:.3f}") |
| print("=" * 64 + "\n") |
|
|
|
|
| def save_metrics_csv(rows, sil, path): |
| import csv |
| with open(path, "w", newline="") as f: |
| w = csv.DictWriter(f, fieldnames=list(rows[0].keys())) |
| w.writeheader() |
| for r in rows: |
| w.writerow(r) |
| f.write(f"# disease_silhouette_real_plus_syn,{sil:.6f}\n") |
| print(f"Saved metrics: {path}") |
|
|
|
|
| |
| def add_kde_contours(ax, emb, diseases, sources, color_map): |
| """Per-disease density contour: real solid, syn dashed (single level).""" |
| from scipy.stats import gaussian_kde |
| xmin, xmax = emb[:, 0].min(), emb[:, 0].max() |
| ymin, ymax = emb[:, 1].min(), emb[:, 1].max() |
| xx, yy = np.mgrid[xmin:xmax:120j, ymin:ymax:120j] |
| grid = np.vstack([xx.ravel(), yy.ravel()]) |
| for d in TARGET_DISEASES: |
| for src, ls in [("real", "-"), ("syn", "--")]: |
| m = (diseases == d) & (sources == src) |
| if m.sum() < 8: |
| continue |
| try: |
| kde = gaussian_kde(emb[m].T) |
| z = kde(grid).reshape(xx.shape) |
| ax.contour(xx, yy, z, levels=[z.max() * 0.30], |
| colors=[color_map[d]], linewidths=0.7, |
| linestyles=ls, alpha=0.55) |
| except Exception: |
| pass |
|
|
|
|
| def plot_overlay(emb, diseases, sources, color_map, axis_label, subtitle, |
| out_prefix, draw_kde=False, syn_marker="x"): |
| fig, ax = plt.subplots(figsize=(5.2, 5.4), dpi=150) |
|
|
| for d in TARGET_DISEASES: |
| c = color_map[d] |
| mr = (diseases == d) & (sources == "real") |
| ms = (diseases == d) & (sources == "syn") |
| |
| ax.scatter(emb[mr, 0], emb[mr, 1], s=15, c=c, marker="o", |
| linewidths=0, alpha=0.55, rasterized=True, zorder=2) |
| |
| if syn_marker == "ring": |
| ax.scatter(emb[ms, 0], emb[ms, 1], s=30, facecolors="none", |
| edgecolors=c, marker="o", linewidths=0.9, |
| alpha=0.9, rasterized=True, zorder=3) |
| else: |
| mk = "x" if syn_marker == "x" else "+" |
| ax.scatter(emb[ms, 0], emb[ms, 1], s=20, c=c, marker=mk, |
| linewidths=0.9, alpha=0.85, rasterized=True, zorder=3) |
|
|
| if draw_kde: |
| add_kde_contours(ax, emb, diseases, sources, color_map) |
|
|
| |
| ax.set_aspect("equal", adjustable="box") |
| xr = emb[:, 0].max() - emb[:, 0].min() |
| yr = emb[:, 1].max() - emb[:, 1].min() |
| pad = 0.06 * max(xr, yr) |
| ax.set_xlim(emb[:, 0].min() - pad, emb[:, 0].max() + pad) |
| ax.set_ylim(emb[:, 1].min() - pad, emb[:, 1].max() + pad) |
|
|
| ax.set_xticks([]); ax.set_yticks([]) |
| ax.set_xlabel(f"{axis_label} 1") |
| ax.set_ylabel(f"{axis_label} 2") |
| if subtitle: |
| ax.set_title(subtitle, fontsize=8, fontweight="normal", color="0.35") |
|
|
| |
| dis_handles = [Line2D([0], [0], marker="o", linestyle="", markersize=5, |
| markerfacecolor=color_map[d], markeredgecolor="none", |
| label=d) for d in TARGET_DISEASES] |
| if syn_marker == "ring": |
| syn_h = Line2D([0], [0], marker="o", linestyle="", markersize=6, |
| markerfacecolor="none", markeredgecolor="0.3", label="Synthetic") |
| else: |
| syn_h = Line2D([0], [0], marker=("x" if syn_marker == "x" else "+"), |
| linestyle="", markersize=6, color="0.3", label="Synthetic") |
| src_handles = [ |
| Line2D([0], [0], marker="o", linestyle="", markersize=5, |
| markerfacecolor="0.3", markeredgecolor="none", label="Real"), |
| syn_h, |
| ] |
| leg1 = fig.legend(handles=dis_handles, loc="lower center", |
| bbox_to_anchor=(0.5, 0.04), ncol=5, |
| columnspacing=1.0, handletextpad=0.3) |
| fig.add_artist(leg1) |
| fig.legend(handles=src_handles, loc="lower center", |
| bbox_to_anchor=(0.5, -0.01), ncol=2, |
| columnspacing=1.5, handletextpad=0.3) |
|
|
| fig.tight_layout(rect=[0, 0.10, 1, 1]) |
| png, svg = f"{out_prefix}.png", f"{out_prefix}.svg" |
| fig.savefig(png, dpi=300, bbox_inches="tight") |
| fig.savefig(svg, bbox_inches="tight") |
| print(f"Saved figure: {png} / {svg}") |
|
|
|
|
| |
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--json", default=JSON_PATH) |
| ap.add_argument("--real_root", default=REAL_ROOT) |
| ap.add_argument("--syn_root", default=SYN_ROOT) |
| ap.add_argument("--out", default=OUT_PREFIX) |
| ap.add_argument("--method", default=EMBED_METHOD, |
| choices=["tsne", "umap", "umap_supervised"]) |
| ap.add_argument("--blur", type=float, default=BLUR_SIGMA) |
| ap.add_argument("--min_dist", type=float, default=UMAP_MIN_DIST) |
| ap.add_argument("--spread", type=float, default=UMAP_SPREAD) |
| ap.add_argument("--n_neighbors", type=int, default=UMAP_NEIGHBORS) |
| ap.add_argument("--syn_marker", default=SYN_MARKER, |
| choices=["x", "plus", "ring"]) |
| ap.add_argument("--kde", action="store_true", default=DRAW_KDE) |
| ap.add_argument("--no_metrics", action="store_true") |
| args = ap.parse_args() |
|
|
| set_nature_style() |
| target_set = set(TARGET_DISEASES) |
|
|
| records = load_records(args.json) |
| print(f"Loaded {len(records)} records from {args.json} " |
| f"(real={REAL_MASK_SOURCE}, syn={SYN_MASK_SOURCE}, " |
| f"method={args.method}, blur={args.blur})") |
|
|
| print("Reading masks...") |
| real_f, real_d = collect(records, args.real_root, "real", target_set, |
| MASK_SIZE, ALLOW_MULTI, REAL_MASK_SOURCE, args.blur) |
| syn_f, syn_d = collect(records, args.syn_root, "syn", target_set, |
| MASK_SIZE, ALLOW_MULTI, SYN_MASK_SOURCE, args.blur) |
|
|
| feats = real_f + syn_f |
| diseases = real_d + syn_d |
| sources = ["real"] * len(real_f) + ["syn"] * len(syn_f) |
| if len(feats) < 10: |
| raise SystemExit("Too few masks loaded — check paths / MASK_SOURCE / disease names.") |
|
|
| feats, diseases, sources = subsample(feats, diseases, sources, |
| MAX_PER_GROUP, RANDOM_STATE) |
|
|
| X = np.stack(feats).astype(np.float32) |
| diseases = np.array(diseases) |
| sources = np.array(sources) |
| print(f"Feature matrix: {X.shape}") |
|
|
| print("Embedding...") |
| emb, X_red = embed(X, diseases, args.method, RANDOM_STATE, |
| n_neighbors=args.n_neighbors, min_dist=args.min_dist, |
| spread=args.spread) |
| color_map = {d: NPG_PALETTE[i] for i, d in enumerate(TARGET_DISEASES)} |
| axis_label = "t-SNE" if args.method == "tsne" else "UMAP" |
|
|
| subtitle = None |
| if COMPUTE_METRICS and not args.no_metrics: |
| print("Computing metrics...") |
| rows, sil = compute_metrics(X_red, diseases, sources, RANDOM_STATE) |
| print_metrics(rows, sil) |
| save_metrics_csv(rows, sil, f"{args.out}_metrics.csv") |
| aucs = [r["c2st_auc"] for r in rows if not np.isnan(r["c2st_auc"])] |
| if aucs: |
| subtitle = f"mean C2ST AUC = {np.mean(aucs):.3f} silhouette = {sil:.3f}" |
|
|
| plot_overlay(emb, diseases, sources, color_map, axis_label, subtitle, |
| args.out, draw_kde=args.kde, syn_marker=args.syn_marker) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
| """ |
| # recommended: UMAP, spread-out clusters, x-markers for syn, metrics |
| python draw_cluster.py \ |
| --json /home/jovyan/AURAD_infer/mask/det-test-cp9000/test_prompt_text2layout_single.json \ |
| --syn_root /home/jovyan/AURAD_infer/mask/det-test-cp9000 \ |
| --real_root /home/jovyan/AURAD_dataset \ |
| --method umap --syn_marker x --min_dist 0.25 --spread 1.0 |
| |
| # if clusters still look too spread / too tight, tune these: |
| python draw_cluster.py ... --min_dist 0.6 --spread 1.5 # looser, fills canvas more |
| python draw_cluster.py ... --min_dist 0.25 --spread 1.0 # tighter clusters |
| |
| # other looks |
| python draw_cluster.py ... --syn_marker ring # open-circle syn (sparse data) |
| python draw_cluster.py ... --method umap_supervised --kde # diseases pushed apart + contours |
| python draw_cluster.py ... --method tsne --blur 0 # original t-SNE style |
| """ |
