| """Multi-arch eval on HER2Match val: cleanfid + pytorch-fid + brown-frac. |
| |
| Supports C7 (UNI-2H encoder), C3 (WaveletJointAEv3), A3 (SkipFreeAE3Stage). |
| Env-var configured for CI-friendly multi-run: |
| ARCH=C7|C3|A3 — which model class to instantiate |
| CKPT_LOCAL — path to best.pt (state_dict must match ARCH) |
| CKPT_TAG — subfolder tag for outputs (e.g. 'c7_best', 'c3_dab_best') |
| OUT_ROOT — where to put generated tiles + results json |
| HF_EXP_SUB — HF sub-path under `augustander/her2match-src-ds-arch-sweep/` for upload |
| """ |
| import os, sys, json, time, glob, subprocess |
| from pathlib import Path |
| import numpy as np |
| import torch, torch.nn.functional as F |
| from PIL import Image |
|
|
| ARCH = os.environ.get('ARCH', 'C7') |
| CKPT_LOCAL = os.environ['CKPT_LOCAL'] |
| CKPT_TAG = os.environ.get('CKPT_TAG', 'best') |
| OUT_ROOT = Path(os.environ.get('OUT_ROOT', f'/home/ubuntu/eval_{ARCH}_{CKPT_TAG}')) |
| HF_EXP_SUB = os.environ.get('HF_EXP_SUB', 'multi_arch_eval') |
| HF_MODEL_REPO = 'augustander/her2match-src-ds-arch-sweep' |
| DATA_ROOT = '/home/ubuntu/her2match_full/her2match_full' |
| GEN_DIR = OUT_ROOT / 'generated_ihc' |
| GEN_DIR.mkdir(parents=True, exist_ok=True) |
| device = 'cuda' |
|
|
| sys.path.insert(0, '/home/ubuntu/work') |
| from src_ds_her2_thunder import ( |
| UniDwtC7, WaveletJointAEv3, SkipFreeAE3Stage, |
| UNI2HSpatialLoss, StainDeconvDAB, |
| ) |
|
|
| |
| def build_model(arch: str, ckpt: dict, uni_loss=None): |
| if arch == 'C7': |
| assert uni_loss is not None, 'C7 needs shared UNI-2H' |
| m = UniDwtC7( |
| uni_module=uni_loss.uni, wavelet='haar', taps=(13, 19), |
| dec_ch=(512, 384, 256, 128, 64), skip_ch=(384, 256, 128), |
| blocks_per_stage=(4, 3, 3, 2, 2), uni_target_size=224, |
| ) |
| elif arch == 'C3': |
| m = WaveletJointAEv3(base_ch=48, wavelet='haar') |
| elif arch == 'A3': |
| m = SkipFreeAE3Stage(base_ch=48, bottleneck_ch=256) |
| else: |
| raise ValueError(f'unknown ARCH={arch}') |
| return m |
|
|
| |
| if ARCH == 'C7': |
| print(f'[{ARCH}/{CKPT_TAG}] loading shared UNI-2H...', flush=True) |
| uni_loss = UNI2HSpatialLoss(target_size=224).to(device).eval() |
| for p in uni_loss.uni.parameters(): p.requires_grad_(False) |
| else: |
| uni_loss = None |
|
|
| |
| G = build_model(ARCH, None, uni_loss=uni_loss).to(device).eval() |
| ckpt = torch.load(CKPT_LOCAL, map_location='cpu', weights_only=False) |
| missing, unexpected = G.load_state_dict(ckpt['model'], strict=False) |
| n_missing_uni = len([k for k in missing if k.startswith('uni.')]) |
| n_missing_other = [k for k in missing if not k.startswith('uni.')] |
| print(f'[{ARCH}/{CKPT_TAG}] loaded ckpt step={ckpt.get("step","?")} ' |
| f'val_lpips={ckpt.get("val_lpips",-1):.4f} ' |
| f'(missing_uni={n_missing_uni}, missing_other={n_missing_other}, unexpected={len(unexpected)})', |
| flush=True) |
| for p in G.parameters(): p.requires_grad_(False) |
| stain_dab = StainDeconvDAB().to(device).eval() |
|
|
| |
| he_val = sorted(glob.glob(f'{DATA_ROOT}/HE/val/*.jpg')) |
| ihc_val_dir = f'{DATA_ROOT}/IHC/val' |
| BATCH = 16 if ARCH == 'C7' else 32 |
| print(f'[{ARCH}/{CKPT_TAG}] generating {len(he_val)} tiles (batch={BATCH})...', flush=True) |
| t0 = time.time() |
| for i in range(0, len(he_val), BATCH): |
| chunk = he_val[i:i+BATCH] |
| imgs = [] |
| for f in chunk: |
| im = Image.open(f).convert('RGB').resize((256, 256), Image.BICUBIC) |
| x = torch.from_numpy(np.asarray(im, dtype=np.float32) / 255.0).permute(2, 0, 1) |
| imgs.append(x * 2 - 1) |
| x = torch.stack(imgs).to(device) |
| with torch.no_grad(): |
| y = G(x).clamp(-1, 1) |
| y_np = ((y.cpu() + 1) * 127.5).clamp(0, 255).byte().permute(0, 2, 3, 1).numpy() |
| for k, f in enumerate(chunk): |
| Image.fromarray(y_np[k]).save(GEN_DIR / Path(f).name, quality=92) |
| if (i // BATCH) % 20 == 0: |
| done = i + len(chunk); rate = done / max(1e-6, time.time() - t0) |
| eta = (len(he_val) - done) / max(1e-6, rate) |
| print(f' gen {done}/{len(he_val)} ({rate:.1f} img/s eta {eta:.0f}s)', flush=True) |
| print(f'[{ARCH}/{CKPT_TAG}] gen done in {time.time()-t0:.0f}s', flush=True) |
|
|
| |
| print(f'[{ARCH}/{CKPT_TAG}] cleanfid...', flush=True) |
| try: |
| from cleanfid import fid as cfid |
| cleanfid_val = cfid.compute_fid( |
| str(GEN_DIR), ihc_val_dir, |
| device='cuda', mode='clean', batch_size=64, num_workers=4) |
| print(f'[{ARCH}/{CKPT_TAG}] CLEANFID = {cleanfid_val:.3f}', flush=True) |
| except Exception as e: |
| print(f'[{ARCH}/{CKPT_TAG}] cleanfid err: {e}', flush=True); cleanfid_val = None |
|
|
| |
| print(f'[{ARCH}/{CKPT_TAG}] pytorch-fid...', flush=True) |
| try: |
| r = subprocess.run( |
| ['python', '-m', 'pytorch_fid', str(GEN_DIR), ihc_val_dir, |
| '--device', 'cuda', '--batch-size', '50'], |
| capture_output=True, text=True, timeout=1800) |
| pytorch_fid_val = None |
| for line in r.stdout.splitlines(): |
| if line.strip().startswith('FID:'): |
| pytorch_fid_val = float(line.split()[1]) |
| break |
| print(f'[{ARCH}/{CKPT_TAG}] PYTORCH_FID = {pytorch_fid_val}', flush=True) |
| if pytorch_fid_val is None: |
| print(f'[{ARCH}/{CKPT_TAG}] pytorch-fid stderr tail: {r.stderr[-500:]}', flush=True) |
| except Exception as e: |
| print(f'[{ARCH}/{CKPT_TAG}] pytorch-fid err: {e}', flush=True); pytorch_fid_val = None |
|
|
| |
| print(f'[{ARCH}/{CKPT_TAG}] brown-frac...', flush=True) |
| def dab_stats(paths, batch=64): |
| means, fracs = [], [] |
| for i in range(0, len(paths), batch): |
| chunk = paths[i:i+batch] |
| imgs = [] |
| for f in chunk: |
| im = Image.open(f).convert('RGB').resize((256, 256), Image.BICUBIC) |
| x = torch.from_numpy(np.asarray(im, dtype=np.float32) / 255.0).permute(2, 0, 1) |
| imgs.append(x * 2 - 1) |
| x = torch.stack(imgs).to(device) |
| with torch.no_grad(): |
| dab = stain_dab(x) |
| means.append(dab.mean(dim=(1,2)).cpu().numpy()) |
| fracs.append((dab > 0.3).float().mean(dim=(1,2)).cpu().numpy()) |
| return np.concatenate(means), np.concatenate(fracs) |
|
|
| gen_files = sorted(GEN_DIR.glob('*.jpg')) |
| tgt_files = sorted(Path(ihc_val_dir).glob('*.jpg')) |
| common = sorted({p.name for p in gen_files} & {p.name for p in tgt_files}) |
| gen_files = [GEN_DIR / n for n in common] |
| tgt_files = [Path(ihc_val_dir) / n for n in common] |
| gm, gf = dab_stats(gen_files) |
| tm, tf = dab_stats(tgt_files) |
| her2p_mask = tf > 0.05 |
| her2p_n = int(her2p_mask.sum()) |
|
|
| results = { |
| 'arch': ARCH, 'ckpt_tag': CKPT_TAG, 'ckpt_local': CKPT_LOCAL, |
| 'ckpt_step': int(ckpt.get('step', -1)), |
| 'ckpt_val_lpips': float(ckpt.get('val_lpips', -1)), |
| 'cleanfid': cleanfid_val, 'pytorch_fid': pytorch_fid_val, |
| 'n_tiles': len(common), 'n_her2p_tiles': her2p_n, |
| 'overall_gen_bf': float(gf.mean()), 'overall_tgt_bf': float(tf.mean()), |
| 'overall_gen_dab': float(gm.mean()), 'overall_tgt_dab': float(tm.mean()), |
| 'her2p_gen_bf': float(gf[her2p_mask].mean()) if her2p_n else None, |
| 'her2p_tgt_bf': float(tf[her2p_mask].mean()) if her2p_n else None, |
| 'her2p_under_predict_pct': float((gf[her2p_mask] < 0.5 * tf[her2p_mask]).mean()) if her2p_n else None, |
| } |
| print(f'[{ARCH}/{CKPT_TAG}] RESULTS:', json.dumps(results, indent=2), flush=True) |
| out_json = OUT_ROOT / f'eval_{ARCH}_{CKPT_TAG}.json' |
| out_json.write_text(json.dumps(results, indent=2)) |
|
|
| if 'HF_TOKEN' not in os.environ and os.path.exists('/tmp/hf_token'): |
| os.environ['HF_TOKEN'] = open('/tmp/hf_token').read().strip() |
| from huggingface_hub import HfApi |
| HfApi().upload_file( |
| path_or_fileobj=str(out_json), |
| path_in_repo=f'{HF_EXP_SUB}/eval_{ARCH}_{CKPT_TAG}.json', |
| repo_id=HF_MODEL_REPO, repo_type='model', |
| commit_message=f'multi-eval: {ARCH} {CKPT_TAG} cleanfid+pytorch-fid+brown-frac') |
| print(f'=== EVAL_DONE {ARCH}/{CKPT_TAG} ===', flush=True) |
|
|
