| """ |
| Evaluate released checkpoints, baselines and trained ablation runs on every split manifest. |
| |
| python -m experiments.evaluate --released --meanmap # released models + input-free baseline |
| python -m experiments.evaluate --released --sweep # + substrate-setting sensitivity |
| python -m experiments.evaluate --stage1 <path>/last_checkpoint_final.pth # image-only stage-1 model |
| python -m experiments.evaluate --runs # every experiments/runs/*/best_checkpoint.pth |
| |
| Metrics are computed PER IMAGE (mean and SD are then taken across images): |
| mse, ms_ssim, pixel_r (Pearson), rme (per-image WFM-RME), mean_gt / mean_pred (heatmap units), |
| mean_gt_Pa / mean_pred_Pa (x colour-scale limit of the set, 500 Pa single cells, 50 Pa spheroids), |
| gt_frac_saturated (pixels >= 0.98 of the colour limit), inference time. |
| Outputs (experiments/results/eval/ by default; existing rows of other models are kept): |
| per_image.csv, summary.csv |
| """ |
| import argparse |
| import glob |
| import json |
| import os |
| import time |
|
|
| import numpy as np |
| import pandas as pd |
| import torch |
|
|
| from experiments.ablation_models import build_generator, n_params |
| from experiments.layout import DEFAULT_DATASET_ROOT, RESULTS_DIR, RUNS_DIR, S2F_ROOT, list_manifests |
| from experiments.losses import PerImageWFMRMELoss |
| from experiments.manifest_data import make_loader |
| from models.s2f_model import create_s2f_model, create_settings_channels |
| from utils.metrics import TorchMetricsWrapper |
| from utils.substrate_settings import compute_settings_normalization |
|
|
| CONFIG_PATH = os.path.join(S2F_ROOT, "config", "substrate_settings.json") |
| SWEEP_KEYS = ["PDMS_1kPa_3um", "fibroblasts_PDMS", "PDMS_10kPa_3um"] |
|
|
|
|
| def released_models(device): |
| out = {} |
| g, _ = create_s2f_model(in_channels=3, model_type="s2f") |
| g.load_state_dict(torch.load(os.path.join(S2F_ROOT, "ckp/single_cell/ckp_singlecell_GN.pth"), map_location="cpu", |
| weights_only=False)["generator_state_dict"], strict=True) |
| out["released_singlecell"] = dict(model=g.to(device).eval(), kind="singlecell", in_channels=3, tanh=True) |
| g, _ = create_s2f_model(in_channels=1, model_type="s2f") |
| g.load_state_dict(torch.load(os.path.join(S2F_ROOT, "ckp/spheroid/ckp_spheroid_GN.pth"), map_location="cpu", |
| weights_only=False)["generator_state_dict"], strict=True) |
| out["released_spheroid_GN"] = dict(model=g.to(device).eval(), kind="spheroid", in_channels=1, tanh=True) |
| fn = os.path.join(S2F_ROOT, "ckp/spheroid/ckp_spheroid_FN.pth") |
| if os.path.exists(fn): |
| g, _ = create_s2f_model(in_channels=1, model_type="s2f_spheroid") |
| g.load_state_dict(torch.load(fn, map_location="cpu", weights_only=False)["model_state_dict"], strict=True) |
| g.set_output_mode(use_tanh=False) |
| out["supervised_spheroid_FN"] = dict(model=g.to(device).eval(), kind="spheroid", in_channels=1, tanh=False) |
| return out |
|
|
|
|
| def stage1_model(path, device): |
| g, _ = create_s2f_model(in_channels=1, model_type="s2f") |
| sd = torch.load(path, map_location="cpu", weights_only=False)["generator_state_dict"] |
| g.load_state_dict({k.replace("denoise_module", "initial_conv"): v for k, v in sd.items()}, strict=True) |
| return dict(model=g.to(device).eval(), kind="singlecell", in_channels=1, tanh=True) |
|
|
|
|
| def run_models(device, pattern): |
| out = {} |
| for p in sorted(glob.glob(pattern)): |
| h = os.path.join(os.path.dirname(p), "history.json") |
| ck = torch.load(p, map_location="cpu", weights_only=False) |
| cfg = ck["config"] |
| done = os.path.exists(h) and len(json.load(open(h))) >= cfg.get("epochs", 0) |
| if not done and not os.environ.get("S2F_EVAL_UNFINISHED"): |
| print(f" skipping unfinished run {os.path.basename(os.path.dirname(p))} (set S2F_EVAL_UNFINISHED=1 to include)") |
| continue |
| g = build_generator(**cfg); g.load_state_dict(ck["generator_state_dict"], strict=True) |
| name = os.path.basename(os.path.dirname(p)) |
| out[name] = dict(model=g.to(device).eval(), kind="singlecell", in_channels=cfg["in_channels"], tanh=True, |
| substrate_mode=cfg["substrate_mode"], input_mode=cfg["input_mode"], |
| epoch=ck.get("epoch"), val_loss=ck.get("val_loss")) |
| return out |
|
|
|
|
| def mean_map(manifest, dataset_root): |
| acc, n = None, 0 |
| for b in make_loader(manifest, dataset_root, batch_size=8, substrate_mode="none"): |
| acc = b[1].sum(0, keepdim=True) if acc is None else acc + b[1].sum(0, keepdim=True); n += b[1].size(0) |
| return acc / n |
|
|
|
|
| def evaluate_one(name, spec, manifest, dataset_root, device, mw, norm, substrate_mode, batch, mean_pred=None): |
| input_mode = spec.get("input_mode", "bf") |
| use_settings = spec.get("in_channels", 1) == 3 |
| sm = substrate_mode if use_settings else "none" |
| forced = None |
| if sm and sm.startswith("force:"): |
| forced, sm = sm.split(":", 1)[1], "reference" |
| try: |
| loader = make_loader(manifest, dataset_root, batch_size=batch, substrate_mode=sm, input_mode=input_mode) |
| except FileNotFoundError as e: |
| print(f" [skip] {name} on {manifest}: {e}"); return [] |
| rows_meta, rec, out, k = loader.dataset.rows, PerImageWFMRMELoss(), [], 0 |
| with torch.no_grad(): |
| for b in loader: |
| x, y = b[0].to(device), b[1].to(device) |
| t0 = time.time() |
| if mean_pred is not None: |
| p = mean_pred.to(device).expand_as(y) |
| else: |
| if use_settings: |
| meta = b[4] |
| if forced: |
| meta = dict(meta, substrate=[forced] * x.size(0)) |
| x = torch.cat([x, create_settings_channels(meta, norm, device, x.shape, config_path=CONFIG_PATH)], 1) |
| p = spec["model"](x) |
| p = (p + 1) / 2 if spec.get("tanh", True) else p |
| dt = (time.time() - t0) / x.size(0) |
| for i in range(y.size(0)): |
| r = rows_meta[k]; k += 1 |
| t, q = y[i:i + 1], p[i:i + 1] |
| yt = float(r.get("ythreshold_Pa") or (50 if "spheroid" in manifest else 500)) |
| out.append({"model": name, "manifest": manifest, "substrate_mode": substrate_mode if use_settings else "n/a", |
| "sample_dir": r["sample_dir"], "set": r.get("set", ""), "condition": r.get("condition", ""), |
| "mse": mw.compute_mse(t, q), "ms_ssim": mw.compute_ms_ssim(t, q), |
| "pixel_r": float(np.corrcoef(t.flatten().cpu().numpy(), q.flatten().cpu().numpy())[0, 1]), |
| "rme": float(rec(q, t)), "mean_gt": float(t.mean()), "mean_pred": float(q.mean()), |
| "mean_gt_Pa": float(t.mean()) * yt, "mean_pred_Pa": float(q.mean()) * yt, |
| "gt_max": float(t.max()), "gt_frac_saturated": float((t >= 0.98).float().mean()), |
| "seconds_per_image": dt}) |
| return out |
|
|
|
|
| def summarise(df): |
| df = df.assign(substrate_mode=df["substrate_mode"].fillna("none")) |
| g = df.groupby(["model", "manifest", "substrate_mode"], dropna=False) |
| s = g.agg(n=("mse", "size"), **{f"{m}_{f}": (m, f) for m in ["mse", "ms_ssim", "pixel_r", "rme", "mean_gt_Pa", |
| "mean_pred_Pa", "gt_frac_saturated"] for f in ["mean", "std"]}, |
| seconds_per_image=("seconds_per_image", "mean")).reset_index() |
| s["pred_over_gt"] = s["mean_pred_Pa_mean"] / s["mean_gt_Pa_mean"] |
| return s |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) |
| ap.add_argument("--released", action="store_true", help="released single-cell + spheroid checkpoints") |
| ap.add_argument("--sweep", action="store_true", help="released single-cell model with 1 / 6 / 10 kPa settings") |
| ap.add_argument("--meanmap", action="store_true", help="input-free baseline (mean training map)") |
| ap.add_argument("--stage1", default=None, help="path to last_checkpoint_final.pth (image-only stage 1)") |
| ap.add_argument("--runs", nargs="?", const=os.path.join(RUNS_DIR, "*", "best_checkpoint.pth"), default=None, |
| help="glob of trained run checkpoints (default experiments/runs/*/best_checkpoint.pth)") |
| ap.add_argument("--manifests", default="", help="comma-separated manifest names (default: all in experiments/splits)") |
| ap.add_argument("--dataset-root", default=DEFAULT_DATASET_ROOT) |
| ap.add_argument("--out", default=os.path.join(RESULTS_DIR, "eval")) |
| ap.add_argument("--batch", type=int, default=2) |
| ap.add_argument("--device", default=None) |
| a = ap.parse_args() |
| device = torch.device(a.device or ("cuda" if torch.cuda.is_available() else "cpu")) |
| mans = [m for m in (a.manifests.split(",") if a.manifests else list_manifests()) if m] |
| sc_mans = [m for m in mans if not m.startswith("spheroid") and m != "alltune_train" and m != "zenodo_v1_train"] |
| sph_mans = [m for m in mans if m.startswith("spheroid") and m != "spheroid_train"] |
| print(f"device {device}; single-cell manifests {sc_mans}; spheroid manifests {sph_mans}") |
| mw, norm = TorchMetricsWrapper(device=device), compute_settings_normalization(config_path=CONFIG_PATH) |
|
|
| models = {} |
| if a.released: |
| models.update(released_models(device)) |
| if a.stage1: |
| models["stage1_image_only"] = stage1_model(a.stage1, device) |
| if a.runs: |
| models.update(run_models(device, a.runs)) |
| rows, params = [], [] |
| for name, spec in models.items(): |
| params.append({"model": name, "n_params_M": n_params(spec["model"]) / 1e6, "epoch": spec.get("epoch"), |
| "val_loss": spec.get("val_loss")}) |
| targets = sc_mans if spec["kind"] == "singlecell" else sph_mans |
| modes = [spec.get("substrate_mode", "legacy")] |
| if name == "released_singlecell": |
| modes = ["legacy", "nominal"] + ([f"force:{k}" for k in SWEEP_KEYS] if a.sweep else []) |
| for m in targets: |
| for sm in modes: |
| print(f" {name:28s} {m:28s} {sm}", flush=True) |
| rows += evaluate_one(name, spec, m, a.dataset_root, device, mw, norm, sm, a.batch) |
| if a.meanmap: |
| for train_m, targets, name in [("alltune_train", sc_mans, "baseline_mean_map_singlecell"), |
| ("spheroid_train", sph_mans, "baseline_mean_map_spheroid")]: |
| if train_m not in list_manifests(): |
| continue |
| M = mean_map(train_m, a.dataset_root) |
| for m in targets: |
| print(f" {name:28s} {m}", flush=True) |
| rows += evaluate_one(name, {"in_channels": 1}, m, a.dataset_root, device, mw, norm, "n/a", a.batch, mean_pred=M) |
| os.makedirs(a.out, exist_ok=True) |
| new = pd.DataFrame(rows) |
| pi = os.path.join(a.out, "per_image.csv") |
| if os.path.exists(pi) and len(new): |
| old = pd.read_csv(pi) |
| new = pd.concat([old[~old.model.isin(new.model.unique())], new], ignore_index=True) |
| if len(new): |
| new.to_csv(pi, index=False) |
| summarise(new).to_csv(os.path.join(a.out, "summary.csv"), index=False) |
| pp = os.path.join(a.out, "model_params.csv") |
| if params: |
| P = pd.DataFrame(params) |
| if os.path.exists(pp): |
| old = pd.read_csv(pp); P = pd.concat([old[~old.model.isin(P.model)], P], ignore_index=True) |
| P.to_csv(pp, index=False) |
| print(f"wrote {pi}, summary.csv, model_params.csv in {a.out}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
