""" On-demand TITAN+CONCH feature extraction for an uploaded whole-slide image. Mirrors the Kaggle extraction notebook but trimmed for a single slide. Heavy on CPU — intended as a best-effort path on free Spaces. Research use only. """ import os, json, gc import numpy as np import h5py PATCH_SIZE = 512 MIN_PATCHES = 64 # lenient on CPU MAX_PATCHES = 1500 # cap to keep CPU runtime sane THUMB_MAX_DIM = 2048 MIN_TISSUE_FRACTION = 0.02 HSV_SAT_MIN, HSV_VAL_MIN, HSV_VAL_MAX = 18, 20, 245 STRIDE_MULTIPLIER = 1.0 _TITAN = None _CONCH = None _TRANSFORM = None def _clean(): gc.collect() try: import torch if torch.cuda.is_available(): torch.cuda.empty_cache() except Exception: pass def load_models(progress=lambda *_: None): """Lazy-load TITAN+CONCH once. Needs HF_TOKEN env var for the gated model.""" global _TITAN, _CONCH, _TRANSFORM if _TITAN is not None: return _TITAN, _CONCH, _TRANSFORM import torch from transformers import AutoModel from huggingface_hub import login from PIL import Image token = os.environ.get("HF_TOKEN", "").strip() if token: login(token=token, add_to_git_credential=False) device = "cuda" if torch.cuda.is_available() else "cpu" progress("Loading TITAN (first run downloads ~GBs)…") titan = AutoModel.from_pretrained("MahmoodLab/TITAN", trust_remote_code=True, token=token or None).to(device).eval() if not hasattr(titan, "return_conch"): raise RuntimeError("TITAN build has no return_conch().") conch, transform = titan.return_conch() conch = conch.to(device).eval() _TITAN, _CONCH, _TRANSFORM = titan, conch, transform progress("Models ready.") return titan, conch, transform def _objective_power(slide): for key in ["openslide.objective-power", "aperio.AppMag", "hamamatsu.SourceLens"]: v = slide.properties.get(key) if v is not None: try: return float(str(v).strip()) except Exception: pass return None def _read_size(slide): obj = _objective_power(slide) if obj is not None and obj >= 35: return 1024 return PATCH_SIZE def _tissue_mask(slide): import cv2 thumb = slide.get_thumbnail((THUMB_MAX_DIM, THUMB_MAX_DIM)).convert("RGB") arr = np.asarray(thumb, np.uint8) hsv = cv2.cvtColor(arr, cv2.COLOR_RGB2HSV) sat, val = hsv[:, :, 1], hsv[:, :, 2] hsv_mask = (sat >= HSV_SAT_MIN) & (val >= HSV_VAL_MIN) & (val <= HSV_VAL_MAX) gray = cv2.cvtColor(arr, cv2.COLOR_RGB2GRAY) _, otsu = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU) mask = (hsv_mask & (otsu > 0)).astype(np.uint8) * 255 kern = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5)) mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kern, 1) mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kern, 2) w0, h0 = slide.dimensions return thumb, mask, w0 / max(thumb.width, 1), h0 / max(thumb.height, 1) def _select_patches(slide): w0, h0 = slide.dimensions read_size = _read_size(slide) thumb, mask, sx, sy = _tissue_mask(slide) max_x, max_y = max(0, w0 - read_size), max(0, h0 - read_size) attempts = [(MIN_TISSUE_FRACTION, STRIDE_MULTIPLIER), (0.02, 0.75), (0.01, 0.75), (0.005, 0.5)] best, best_info = [], None for min_tissue, stride_mult in attempts: stride = max(1, int(round(read_size * stride_mult))) cand = [] for y in range(0, max_y + 1, stride): ty0 = max(0, min(mask.shape[0] - 1, int(y / sy))) ty1 = max(ty0 + 1, min(mask.shape[0], int((y + read_size) / sy))) for x in range(0, max_x + 1, stride): tx0 = max(0, min(mask.shape[1] - 1, int(x / sx))) tx1 = max(tx0 + 1, min(mask.shape[1], int((x + read_size) / sx))) if float(mask[ty0:ty1, tx0:tx1].mean() / 255.0) >= min_tissue: cand.append((int(x), int(y))) if len(cand) > len(best): best, best_info = cand, (min_tissue, stride_mult, stride) if len(cand) >= MIN_PATCHES: break if len(best) < MIN_PATCHES: raise RuntimeError(f"Only {len(best)} tissue patches found (need {MIN_PATCHES}).") best = sorted(best, key=lambda t: (t[1], t[0])) if len(best) > MAX_PATCHES: idx = np.linspace(0, len(best) - 1, MAX_PATCHES).astype(int) best = [best[i] for i in idx] coords = np.array(best, dtype=np.int64) return coords, int(read_size) def extract_to_h5(wsi_path, out_h5, progress=lambda *_: None, batch_size=8): """Extract TITAN+CONCH features from a whole slide → write an .h5 matching the schema.""" import torch import torch.nn.functional as F import openslide from PIL import Image titan, conch, transform = load_models(progress) device = "cuda" if torch.cuda.is_available() else "cpu" progress("Opening slide…") slide = openslide.OpenSlide(wsi_path) try: progress("Selecting tissue patches…") coords, read_size = _select_patches(slide) n = len(coords) progress(f"Encoding {n} patches with CONCH (CPU is slow)…") feats = [] for start in range(0, n, batch_size): chunk = coords[start:start + batch_size] imgs = [] for (x, y) in chunk: patch = slide.read_region((int(x), int(y)), 0, (read_size, read_size)).convert("RGB") if read_size != PATCH_SIZE: patch = patch.resize((PATCH_SIZE, PATCH_SIZE), Image.BILINEAR) imgs.append(transform(patch)) batch = torch.stack(imgs, 0).to(device) with torch.inference_mode(): out = conch.encode_image(batch) if hasattr(conch, "encode_image") else conch(batch) if isinstance(out, (tuple, list)): out = out[0] if out.ndim == 3: out = out[:, 0, :] feats.append(F.normalize(out.float().cpu(), dim=1).numpy().astype(np.float32)) progress(f"Encoded {min(start + batch_size, n)}/{n} patches…") _clean() features = np.concatenate(feats, 0).astype(np.float32) progress("Computing TITAN slide embedding…") ft = torch.from_numpy(features).float().to(device) ct = torch.from_numpy(coords).long().to(device) with torch.inference_mode(): emb = titan.encode_slide_from_patch_features(ft, ct, int(read_size)) if isinstance(emb, (tuple, list)): emb = emb[0] if emb.ndim == 2 and emb.shape[0] == 1: emb = emb[0] titan_emb = emb.reshape(-1).float().cpu().numpy().astype(np.float32) nrm = float(np.linalg.norm(titan_emb)) if nrm > 1e-8: titan_emb /= nrm pooled = np.concatenate([features.mean(0), features.max(0), features.std(0)]).astype(np.float32) pn = float(np.linalg.norm(pooled)) if pn > 1e-8: pooled /= pn progress("Writing .h5…") sid = os.path.splitext(os.path.basename(wsi_path))[0] with h5py.File(out_h5, "w") as h5: h5.create_dataset("features", data=features, compression="gzip", compression_opts=4) h5.create_dataset("coords", data=coords, compression="gzip", compression_opts=4) h5.create_dataset("titan_slide_embedding", data=titan_emb, compression="gzip", compression_opts=4) h5.create_dataset("pooled_conch_slide_feature_mean_max_std", data=pooled, compression="gzip", compression_opts=4) h5["coords"].attrs["patch_size_level0"] = int(read_size) h5.attrs["slide_id"] = sid h5.attrs["label_binary"] = -1 h5.attrs["subtype"] = "TEST" h5.attrs["meta_json"] = json.dumps({"read_size_level0": int(read_size), "num_patches": int(n)}) progress("Extraction complete.") return out_h5 finally: slide.close() _clean()