| """GigaTIME (Providence/Microsoft, Cell 2025) virtual mIF inference. |
| |
| Model : prov-gigatime/GigaTIME (HF, gated, Apache-2.0) -- UNet++ CNN, H&E RGB -> 23 sigmoid maps |
| (21 proteins + 2 background channels TRITC/Cy5, which are dropped here). |
| Code : https://github.com/prov-gigatime/GigaTIME (architecture re-implemented below, verbatim |
| layer layout of scripts/archs.py, so nothing from the downloaded repo is imported). |
| |
| Preprocessing (as in the official notebooks): RGB/255, ImageNet mean/std normalisation, no stain norm. |
| Optional bilinear upsampling by --scale before inference (the model was trained on native-resolution |
| Providence tiles; ORION/PathGen tiles are ~0.5 um/px, so --scale 2 brings them to ~0.25 um/px = 40x). |
| The net is fully convolutional (4 poolings): images are reflect-padded to a multiple of 16 and run |
| whole; images larger than --max_full px are run as 512-px tiles with a 64-px overlap margin |
| (only the tile centres are kept), which avoids the seams of the official non-overlapping 256 windows. |
| Outputs are sigmoid probabilities (the official code thresholds at 0.5; we keep the probabilities). |
| |
| Usage (HF token only via environment, never stored): |
| HF_TOKEN=... .venv/bin/python 25_gigatime_infer.py --set orion --scale 2 |
| .venv/bin/python 25_gigatime_infer.py --set orion --scale 1 --subset 200 --tag scaletest_s1 --no_maps |
| .venv/bin/python 25_gigatime_infer.py --set tcga_text --scale 2 |
| .venv/bin/python 25_gigatime_infer.py --set ge_patches --scale 2 |
| |
| Outputs (cache/gigatime/): |
| <set>_maps.npy uint8 (N, 21, R, R) probability*255, area-resampled to R=--map_res (default 128) |
| <set>_means.npy float32 (N, 21) per-tile channel mean probability (at inference resolution, |
| padding excluded) |
| <set>_frac.npy float32 (N, 21) per-tile fraction of pixels with prob > 0.5 |
| <set>_meta.csv row order, source path, status, sizes |
| channels.json |
| orion: orion_validation.csv/json (shared-channel agreement with the real ORION mIF) |
| """ |
| import argparse |
| import glob |
| import json |
| import os |
| import queue |
| import sys |
| import threading |
| import time |
|
|
| import numpy as np |
| import pandas as pd |
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| from PIL import Image, PngImagePlugin |
|
|
| PngImagePlugin.MAX_TEXT_CHUNK = 64 * 1024 * 1024 |
|
|
| ROOT = os.path.dirname(os.path.abspath(__file__)) |
| sys.path.insert(0, ROOT) |
| from common import CACHE, ORION, load_all_tiles |
|
|
| OUT = os.path.join(CACHE, "gigatime") |
| WDIR = os.path.join(ROOT, "ext", "gigatime", "weights", "GigaTIME") |
| PATHGEN = os.path.join(ROOT, "..", "PathGen-1.6M") |
| GE_DIRS = ["brain_image", "brain_train_image", "brain_trainnew_image", "gliomas_test_image_new", |
| "gliomas_train_image", "idc_image", "idc_train_image", "saromas_test_image_new", |
| "saromas_train_image_new", "skin_image", "skin_train_image"] |
|
|
| ALL_CHANNELS = ['DAPI', 'TRITC', 'Cy5', 'PD-1', 'CD14', 'CD4', 'T-bet', 'CD34', 'CD68', 'CD16', 'CD11c', |
| 'CD138', 'CD20', 'CD3', 'CD8', 'PD-L1', 'CK', 'Ki67', 'Tryptase', 'Actin-D', 'Caspase3-D', |
| 'PHH3-B', 'Transgelin'] |
| BACKGROUND = ['TRITC', 'Cy5'] |
| KEEP = [i for i, c in enumerate(ALL_CHANNELS) if c not in BACKGROUND] |
| CHANNELS = [ALL_CHANNELS[i] for i in KEEP] |
| |
| SHARED = {'DAPI': 'Hoechst', 'PD-1': 'PD-1', 'CD4': 'CD4', 'CD68': 'CD68', 'CD20': 'CD20', 'CD3': 'CD3e', |
| 'CD8': 'CD8a', 'PD-L1': 'PD-L1', 'CK': 'Pan-CK', 'Ki67': 'Ki67', |
| |
| 'Actin-D': 'SMA', 'Transgelin': 'SMA'} |
| MEAN = torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1) |
| STD = torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1) |
|
|
|
|
| |
| class VGGBlock(nn.Module): |
| def __init__(self, i, m, o): |
| super().__init__() |
| self.relu = nn.ReLU(inplace=True) |
| self.conv1, self.bn1 = nn.Conv2d(i, m, 3, padding=1), nn.BatchNorm2d(m) |
| self.conv2, self.bn2 = nn.Conv2d(m, o, 3, padding=1), nn.BatchNorm2d(o) |
|
|
| def forward(self, x): |
| return self.relu(self.bn2(self.conv2(self.relu(self.bn1(self.conv1(x)))))) |
|
|
|
|
| class GigaTIME(nn.Module): |
| def __init__(self, num_classes=23, input_channels=3): |
| super().__init__() |
| f = [32, 64, 128, 256, 512] |
| self.pool = nn.MaxPool2d(2, 2) |
| self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True) |
| self.conv0_0 = VGGBlock(input_channels, f[0], f[0]) |
| self.conv1_0 = VGGBlock(f[0], f[1], f[1]) |
| self.conv2_0 = VGGBlock(f[1], f[2], f[2]) |
| self.conv3_0 = VGGBlock(f[2], f[3], f[3]) |
| self.conv4_0 = VGGBlock(f[3], f[4], f[4]) |
| self.conv0_1 = VGGBlock(f[0] + f[1], f[0], f[0]) |
| self.conv1_1 = VGGBlock(f[1] + f[2], f[1], f[1]) |
| self.conv2_1 = VGGBlock(f[2] + f[3], f[2], f[2]) |
| self.conv3_1 = VGGBlock(f[3] + f[4], f[3], f[3]) |
| self.conv0_2 = VGGBlock(f[0] * 2 + f[1], f[0], f[0]) |
| self.conv1_2 = VGGBlock(f[1] * 2 + f[2], f[1], f[1]) |
| self.conv2_2 = VGGBlock(f[2] * 2 + f[3], f[2], f[2]) |
| self.conv0_3 = VGGBlock(f[0] * 3 + f[1], f[0], f[0]) |
| self.conv1_3 = VGGBlock(f[1] * 3 + f[2], f[1], f[1]) |
| self.conv0_4 = VGGBlock(f[0] * 4 + f[1], f[0], f[0]) |
| self.final = nn.Conv2d(f[0], num_classes, kernel_size=1) |
|
|
| def forward(self, x): |
| up, pool = self.up, self.pool |
| x0_0 = self.conv0_0(x) |
| x1_0 = self.conv1_0(pool(x0_0)) |
| x0_1 = self.conv0_1(torch.cat([x0_0, up(x1_0)], 1)) |
| x2_0 = self.conv2_0(pool(x1_0)) |
| x1_1 = self.conv1_1(torch.cat([x1_0, up(x2_0)], 1)) |
| x0_2 = self.conv0_2(torch.cat([x0_0, x0_1, up(x1_1)], 1)) |
| x3_0 = self.conv3_0(pool(x2_0)) |
| x2_1 = self.conv2_1(torch.cat([x2_0, up(x3_0)], 1)) |
| x1_2 = self.conv1_2(torch.cat([x1_0, x1_1, up(x2_1)], 1)) |
| x0_3 = self.conv0_3(torch.cat([x0_0, x0_1, x0_2, up(x1_2)], 1)) |
| x4_0 = self.conv4_0(pool(x3_0)) |
| x3_1 = self.conv3_1(torch.cat([x3_0, up(x4_0)], 1)) |
| x2_2 = self.conv2_2(torch.cat([x2_0, x2_1, up(x3_1)], 1)) |
| x1_3 = self.conv1_3(torch.cat([x1_0, x1_1, x1_2, up(x2_2)], 1)) |
| x0_4 = self.conv0_4(torch.cat([x0_0, x0_1, x0_2, x0_3, up(x1_3)], 1)) |
| return self.final(x0_4) |
|
|
|
|
| def load_model(device): |
| wp = os.path.join(WDIR, "model.pth") |
| if not os.path.exists(wp): |
| from huggingface_hub import snapshot_download |
| snapshot_download("prov-gigatime/GigaTIME", local_dir=WDIR) |
| sd = torch.load(wp, map_location="cpu", weights_only=True) |
| m = GigaTIME(23, 3) |
| m.load_state_dict(sd, strict=True) |
| return m.eval().to(device) |
|
|
|
|
| |
| def _pad_reflect(x, top, bottom, left, right): |
| |
| mode = "reflect" if max(top, bottom) < x.shape[2] and max(left, right) < x.shape[3] else "replicate" |
| return F.pad(x, (left, right, top, bottom), mode=mode) |
|
|
|
|
| @torch.no_grad() |
| def infer(model, x, max_full=768, tile=512, margin=64, tile_bs=4): |
| """x: (1,3,H,W) normalised float on device -> logits (1,23,H,W) float32.""" |
| _, _, H, W = x.shape |
| with torch.autocast("cuda", dtype=torch.float16): |
| if max(H, W) <= max_full: |
| xp = _pad_reflect(x, 0, (-H) % 16, 0, (-W) % 16) |
| return model(xp)[:, :, :H, :W].float() |
| step = tile - 2 * margin |
| ny, nx = -(-H // step), -(-W // step) |
| xp = _pad_reflect(x, margin, ny * step + margin - H, margin, nx * step + margin - W) |
| out = torch.zeros(1, 23, ny * step, nx * step, device=x.device) |
| pos = [(iy, ix) for iy in range(ny) for ix in range(nx)] |
| for k in range(0, len(pos), tile_bs): |
| chunk = pos[k:k + tile_bs] |
| tiles = torch.cat([xp[:, :, iy * step:iy * step + tile, ix * step:ix * step + tile] for iy, ix in chunk]) |
| o = model(tiles)[:, :, margin:margin + step, margin:margin + step].float() |
| for j, (iy, ix) in enumerate(chunk): |
| out[:, :, iy * step:(iy + 1) * step, ix * step:(ix + 1) * step] = o[j] |
| return out[:, :, :H, :W] |
|
|
|
|
| def read_rgb(path): |
| with Image.open(path) as im: |
| return np.asarray(im.convert("RGB")) |
|
|
|
|
| def prefetch(paths, extra=None, qsize=16): |
| q = queue.Queue(qsize) |
|
|
| def work(): |
| for i, p in enumerate(paths): |
| try: |
| img = read_rgb(p) |
| ex = extra(i) if extra else None |
| q.put((i, img, ex, None)) |
| except Exception as e: |
| q.put((i, None, None, f"{type(e).__name__}: {str(e)[:200]}")) |
| q.put(None) |
|
|
| threading.Thread(target=work, daemon=True).start() |
| while True: |
| it = q.get() |
| if it is None: |
| return |
| yield it |
|
|
|
|
| |
| def get_set(name, subset, seed=0): |
| if name == "orion": |
| df = load_all_tiles() |
| meta = pd.DataFrame({"tile_id": df.tile_id, "source_path": [os.path.normpath(os.path.join(ORION, p)) for p in df.image_path], |
| "if_path": [os.path.normpath(os.path.join(ORION, p)) for p in df.target_path], |
| "orion_split": df.orion_split}) |
| prop = df[[c + "_prop" for c in sorted(set(SHARED.values()), key=list(SHARED.values()).index)]] |
| elif name == "tcga_text": |
| d = pd.read_csv(os.path.join(CACHE, "allimg_meta.csv")) |
| meta = pd.DataFrame({"idx": d["idx"], "key": d["key"], "source_path": d.image_path.map(os.path.normpath)}) |
| prop = None |
| elif name == "ge_patches": |
| files = [] |
| for g in GE_DIRS: |
| files += [(g, f) for f in sorted(glob.glob(os.path.join(PATHGEN, g, "*"))) if os.path.isfile(f) and f.lower().endswith((".png", ".jpg", ".jpeg", ".tif", ".tiff"))] |
| meta = pd.DataFrame({"group": [g for g, _ in files], "source_path": [os.path.normpath(f) for _, f in files]}) |
| prop = None |
| else: |
| raise ValueError(name) |
| if subset: |
| sel = np.sort(np.random.default_rng(seed).choice(len(meta), subset, replace=False)) |
| meta = meta.iloc[sel].reset_index(drop=True) |
| prop = prop.iloc[sel].reset_index(drop=True) if prop is not None else None |
| return meta.reset_index(drop=True), prop |
|
|
|
|
| def pearson_rows(a, b): |
| """row-wise Pearson for (n, p) arrays; NaN where a row is constant.""" |
| a = a - a.mean(1, keepdims=True) |
| b = b - b.mean(1, keepdims=True) |
| den = np.sqrt((a * a).sum(1) * (b * b).sum(1)) |
| with np.errstate(invalid="ignore", divide="ignore"): |
| return np.where(den > 0, (a * b).sum(1) / den, np.nan) |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--set", required=True, choices=["orion", "tcga_text", "ge_patches"]) |
| ap.add_argument("--scale", type=float, default=2.0) |
| ap.add_argument("--subset", type=int, default=0) |
| ap.add_argument("--tag", default="") |
| ap.add_argument("--no_maps", action="store_true") |
| ap.add_argument("--map_res", type=int, default=128) |
| ap.add_argument("--pix_tiles", type=int, default=400, help="ORION tiles used for pixel-level Pearson") |
| ap.add_argument("--gpu_frac", type=float, default=0.2, help="cap on this process' share of GPU memory") |
| a = ap.parse_args() |
|
|
| os.makedirs(OUT, exist_ok=True) |
| dev = "cuda" |
| torch.cuda.set_per_process_memory_fraction(a.gpu_frac) |
| torch.backends.cudnn.benchmark = False |
| model = load_model(dev) |
| mean, std = MEAN.to(dev), STD.to(dev) |
| tag = a.tag or a.set |
| meta, prop = get_set(a.set, a.subset) |
| N, C, R = len(meta), len(KEEP), a.map_res |
| print(f"[{tag}] N={N} scale={a.scale} channels={C}", flush=True) |
|
|
| with open(os.path.join(OUT, "channels.json"), "w") as f: |
| json.dump({"channels": CHANNELS, "model_output_channels": ALL_CHANNELS, "dropped_background": BACKGROUND, |
| "orion_shared": SHARED, "value": "sigmoid probability (maps: uint8 = round(prob*255))", |
| "source": "huggingface.co/prov-gigatime/GigaTIME"}, f, indent=1) |
|
|
| maps = None if a.no_maps else np.lib.format.open_memmap(os.path.join(OUT, f"{tag}_maps.npy"), "w+", np.uint8, (N, C, R, R)) |
| means = np.full((N, C), np.nan, np.float32) |
| frac = np.full((N, C), np.nan, np.float32) |
| status, hw, inhw = ["failed: not processed"] * N, [(0, 0)] * N, [(0, 0)] * N |
|
|
| is_orion = a.set == "orion" |
| if is_orion: |
| import tifffile |
| orion_ch = ["Hoechst", "CD31", "CD45", "CD68", "CD4", "FOXP3", "CD8a", "CD45RO", "CD20", "PD-L1", "CD3e", |
| "CD163", "E-cadherin", "PD-1", "Ki67", "Pan-CK", "SMA"] |
| pairs = [(g, CHANNELS.index(g), o, orion_ch.index(o)) for g, o in SHARED.items()] |
| real_mean = np.full((N, len(orion_ch)), np.nan, np.float32) |
| pix_sel = set(np.linspace(0, N - 1, min(a.pix_tiles, N)).astype(int).tolist()) |
| pix_r = {g: [] for g, *_ in pairs} |
| box_r = {g: [] for g, *_ in pairs} |
| pooled = {g: ([], []) for g, *_ in pairs} |
| rng = np.random.default_rng(0) |
| extra = lambda i: tifffile.imread(meta.if_path[i]) |
| else: |
| extra = None |
|
|
| t0 = time.time() |
| for i, img, ifa, err in prefetch(list(meta.source_path), extra): |
| if err is not None: |
| status[i] = "failed: " + err |
| print(f" FAIL {i} {meta.source_path[i]} {err}", flush=True) |
| continue |
| try: |
| H, W = img.shape[:2] |
| x = torch.from_numpy(img).to(dev).permute(2, 0, 1)[None].float() / 255.0 |
| if a.scale != 1: |
| x = F.interpolate(x, scale_factor=a.scale, mode="bilinear", align_corners=False).clamp(0, 1) |
| x = (x - mean) / std |
| p = torch.sigmoid(infer(model, x))[:, KEEP] |
| means[i] = p.mean((0, 2, 3)).cpu().numpy() |
| frac[i] = (p > 0.5).float().mean((0, 2, 3)).cpu().numpy() |
| if maps is not None: |
| maps[i] = (F.interpolate(p, size=(R, R), mode="area")[0] * 255).round().clamp(0, 255).byte().cpu().numpy() |
| status[i], hw[i], inhw[i] = "ok", (H, W), tuple(p.shape[2:]) |
| if is_orion: |
| real_mean[i] = ifa.reshape(-1, ifa.shape[2]).mean(0) |
| if i in pix_sel: |
| pr = F.interpolate(p, size=(H, W), mode="area")[0].cpu().numpy() |
| for g, gi, o, oi in pairs: |
| a1, b1 = pr[gi].ravel()[None], ifa[..., oi].astype(np.float32).ravel()[None] |
| pix_r[g].append(pearson_rows(a1, b1)[0]) |
| bh, bw = H // 8 * 8, W // 8 * 8 |
| a8 = pr[gi][:bh, :bw].reshape(bh // 8, 8, bw // 8, 8).mean((1, 3)).ravel()[None] |
| b8 = ifa[:bh, :bw, oi].astype(np.float32).reshape(bh // 8, 8, bw // 8, 8).mean((1, 3)).ravel()[None] |
| box_r[g].append(pearson_rows(a8, b8)[0]) |
| s = rng.choice(a1.shape[1], 2000, replace=False) |
| pooled[g][0].append(a1[0, s]); pooled[g][1].append(b1[0, s]) |
| except Exception as e: |
| status[i] = f"failed: {type(e).__name__}: {str(e)[:200]}" |
| print(f" FAIL {i} {meta.source_path[i]} {status[i]}", flush=True) |
| if (i + 1) % 500 == 0: |
| print(f" {i + 1}/{N} {time.time() - t0:.0f}s gpu_max={torch.cuda.max_memory_allocated() / 2**30:.2f}GB", flush=True) |
|
|
| if maps is not None: |
| maps.flush() |
| meta["status"] = status |
| meta["orig_h"], meta["orig_w"] = [h for h, _ in hw], [w for _, w in hw] |
| meta["infer_h"], meta["infer_w"] = [h for h, _ in inhw], [w for _, w in inhw] |
| meta["scale"] = a.scale |
| meta.insert(0, "row", np.arange(N)) |
| meta.to_csv(os.path.join(OUT, f"{tag}_meta.csv"), index=False) |
| np.save(os.path.join(OUT, f"{tag}_means.npy"), means) |
| np.save(os.path.join(OUT, f"{tag}_frac.npy"), frac) |
| nok = int((meta.status == "ok").sum()) |
| print(f"[{tag}] done ok={nok} failed={N - nok} time={time.time() - t0:.0f}s " |
| f"gpu_max={torch.cuda.max_memory_allocated() / 2**30:.2f}GB", flush=True) |
|
|
| if is_orion: |
| from scipy.stats import spearmanr, pearsonr |
| ok = (meta.status == "ok").values |
| rows = [] |
| for g, gi, o, oi in pairs: |
| r = dict(gigatime=g, orion=o, |
| tile_mean_spearman=spearmanr(means[ok, gi], real_mean[ok, oi])[0], |
| tile_frac_spearman=spearmanr(frac[ok, gi], real_mean[ok, oi])[0], |
| tile_mean_vs_prop_spearman=spearmanr(means[ok, gi], prop[o + "_prop"].values[ok])[0], |
| pixel_pearson_tile_median=float(np.nanmedian(pix_r[g])), |
| pixel_pearson_tile_mean=float(np.nanmean(pix_r[g])), |
| pixel_pearson_pooled=pearsonr(np.concatenate(pooled[g][0]), np.concatenate(pooled[g][1]))[0], |
| box8_pearson_tile_mean=float(np.nanmean(box_r[g])), |
| n_tiles=int(ok.sum()), n_pix_tiles=len(pix_r[g])) |
| rows.append(r) |
| v = pd.DataFrame(rows) |
| v.to_csv(os.path.join(OUT, f"{tag}_validation.csv"), index=False) |
| core = v[~v.gigatime.isin(["Actin-D", "Transgelin"])] |
| summ = {"scale": a.scale, "N": int(ok.sum()), |
| "mean_over_10_core_channels": core.drop(columns=["gigatime", "orion"]).mean().to_dict()} |
| with open(os.path.join(OUT, f"{tag}_validation.json"), "w") as f: |
| json.dump(summ, f, indent=1, default=float) |
| np.save(os.path.join(OUT, f"{tag}_real_if_means.npy"), real_mean) |
| pd.set_option("display.width", 250) |
| print(v.round(3).to_string(index=False)) |
| print(json.dumps(summ, indent=1, default=float)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
