File size: 11,755 Bytes
1d4bac8 3c48411 1d4bac8 5004189 1d4bac8 5004189 1d4bac8 5004189 1d4bac8 3c48411 1d4bac8 f557147 1d4bac8 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 | """
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")) # "n/a" is read back from CSV as NaN
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): # keep rows of models not re-evaluated now
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()
|