Spaces:
Running
Running
| """ | |
| 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() | |