"""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/): _maps.npy uint8 (N, 21, R, R) probability*255, area-resampled to R=--map_res (default 128) _means.npy float32 (N, 21) per-tile channel mean probability (at inference resolution, padding excluded) _frac.npy float32 (N, 21) per-tile fraction of pixels with prob > 0.5 _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 # noqa: E402 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] # GigaTIME channel -> ORION channel (ORION if/*.tiff channel order == *_prop column order) 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', # tentative (smooth-muscle markers vs SMA) '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) # ---------------- architecture (same layout as GigaTIME scripts/archs.py, UNet++) ---------------- 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): # token is read from the HF_TOKEN env var by huggingface_hub 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) # ---------------- inference ---------------- def _pad_reflect(x, top, bottom, left, right): # reflect padding must be < dim; fall back to replicate for tiny images 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: # unreadable -> logged as failed 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 # ---------------- datasets ---------------- 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]) # (333,333,17) uint8 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] # (1,C,h,w) 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() # back onto the IF grid 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 # 8x8-pixel boxes (~4 um), as in the official eval 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()