"""Deep feature extraction backbones for image clustering. The primary backbone is EnFormer (Deep Ensemble Clustering backbone, ImageNet-1K pretrained), used here purely as a frozen feature extractor: an image is mapped to the 320-d representation produced just before the classification head. A torchvision ConvNeXt-Tiny backbone is provided as a fallback so the app still runs if the EnFormer weights cannot be fetched. """ import os import sys import urllib.request from typing import List, Optional, Tuple import numpy as np import torch import torch.nn.functional as F from PIL import Image from scipy import ndimage # Make the vendored EnFormer model package importable. _HERE = os.path.dirname(os.path.abspath(__file__)) if _HERE not in sys.path: sys.path.insert(0, _HERE) IMAGENET_MEAN = torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1) IMAGENET_STD = torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1) ENFORMER_SMALL_URL = ( "https://github.com/En-Former/EnFormer/releases/download/v1.0/enformer_small.pth.tar" ) def _default_weight_path() -> str: cache = os.environ.get("ENFORMER_WEIGHTS_DIR") or os.path.join( os.path.expanduser("~"), ".cache", "enformer_cluster" ) os.makedirs(cache, exist_ok=True) return os.path.join(cache, "enformer_small.pth.tar") def _resolve_local_weight() -> Optional[str]: """Look for a bundled checkpoint next to the app before downloading.""" candidates = [ os.path.join(_HERE, "weights", "enformer_small.pth"), os.path.join(_HERE, "weights", "enformer_small.pth.tar"), os.path.join(_HERE, "enformer_small.pth"), os.path.join(os.path.dirname(_HERE), "weights", "enformer_small.pth.tar"), ] for c in candidates: if os.path.isfile(c): return c return None def download_enformer_weights(dst: Optional[str] = None) -> str: dst = dst or _default_weight_path() if os.path.isfile(dst) and os.path.getsize(dst) > 1_000_000: return dst tmp = dst + ".part" urllib.request.urlretrieve(ENFORMER_SMALL_URL, tmp) os.replace(tmp, dst) return dst class FeatureExtractor: """Frozen deep feature extractor. Parameters ---------- backbone : {"enformer", "convnext"} Which pretrained network to use. "enformer" is preferred; "convnext" is a self-contained torchvision fallback. device : str weights_path : optional explicit path to the EnFormer checkpoint. """ def __init__(self, backbone: str = "enformer", device: Optional[str] = None, weights_path: Optional[str] = None): self.device = device or ("cuda" if torch.cuda.is_available() else "cpu") self.backbone_name = backbone self.feature_dim: int if backbone == "enformer": self._build_enformer(weights_path) elif backbone == "convnext": self._build_convnext() else: raise ValueError(f"unknown backbone {backbone!r}") # NB: EnsembleFormer overrides train() and returns None, so eval() also # returns None -- do not chain .eval().to(...). self.model.to(self.device) self.model.eval() # ------------------------------------------------------------------ builders def _build_enformer(self, weights_path: Optional[str]): # The vendored model prints optional-import notices (mmdet/mmseg) at import # time; silence them for a clean log. import contextlib import io as _io with contextlib.redirect_stdout(_io.StringIO()): from enformer_models import enformer_small model = enformer_small(num_classes=1000) path = weights_path or _resolve_local_weight() if path is None: path = download_enformer_weights() ckpt = torch.load(path, map_location="cpu", weights_only=False) sd = ckpt.get("state_dict", ckpt.get("model", ckpt)) missing, unexpected = model.load_state_dict(sd, strict=False) # head is unused for feature extraction; tolerate only that mismatch real_missing = [k for k in missing if not k.startswith("head")] if real_missing: raise RuntimeError(f"EnFormer weights missing keys: {real_missing[:6]}") self.model = model self.feature_dim = 320 def _build_convnext(self): import torchvision weights = torchvision.models.ConvNeXt_Tiny_Weights.IMAGENET1K_V1 net = torchvision.models.convnext_tiny(weights=weights) net.classifier = torch.nn.Flatten(1) # keep pooled 768-d feature self.model = net self.feature_dim = 768 # ------------------------------------------------------------------ forward def _forward_enformer(self, x: torch.Tensor) -> torch.Tensor: m = self.model x = m.forward_embeddings(x) for stage in m.stages: x, _ = stage(x) return m.classifier_pre_norm(x.mean([-2, -1])) # [B, 320] def _forward(self, x: torch.Tensor) -> torch.Tensor: if self.backbone_name == "enformer": return self._forward_enformer(x) return self.model(x) # ------------------------------------------------------------------ preprocess @staticmethod def _to_tensor(img: Image.Image, resolution: int) -> torch.Tensor: img = img.convert("RGB").resize((resolution, resolution), Image.BILINEAR) arr = torch.from_numpy(np.asarray(img).copy()).float().permute(2, 0, 1) / 255.0 return arr def _preprocess_batch(self, imgs: List[Image.Image], resolution: int, tiles: int) -> torch.Tensor: """Return a batch tensor. When tiles>1 each image is split into a tiles x tiles grid of crops (each resized to `resolution`); crop features are averaged later.""" if tiles <= 1: batch = torch.stack([self._to_tensor(im, resolution) for im in imgs]) else: crops = [] for im in imgs: im = im.convert("RGB") w, h = im.size cw, ch = w // tiles, h // tiles for i in range(tiles): for j in range(tiles): box = (j * cw, i * ch, (j + 1) * cw, (i + 1) * ch) crops.append(self._to_tensor(im.crop(box), resolution)) batch = torch.stack(crops) batch = (batch - IMAGENET_MEAN) / IMAGENET_STD return batch # ------------------------------------------------------------------ public API @torch.no_grad() def extract(self, images: List, resolution: int = 224, tiles: int = 1, batch_size: int = 16, l2_normalize: bool = True, progress=None) -> np.ndarray: """Extract one feature vector per input image. images : list of PIL.Image or file paths. tiles : split each image into tiles x tiles crops and average their features (captures fine texture on high-resolution micrographs). """ # EnFormer's PartitionalClustering uses gumbel_softmax, which samples noise # even in eval mode. Fix the RNG so identical inputs give identical features # (reproducible clustering). Save/restore state to avoid side effects. rng_state = torch.get_rng_state() torch.manual_seed(0) feats = [] n = len(images) crops_per_img = tiles * tiles # process image-by-image group so tiling stays aligned img_bs = max(1, batch_size // crops_per_img) for start in range(0, n, img_bs): chunk = images[start:start + img_bs] pil = [Image.open(x) if isinstance(x, str) else x for x in chunk] batch = self._preprocess_batch(pil, resolution, tiles).to(self.device) out = self._forward(batch) if tiles > 1: out = out.view(len(pil), crops_per_img, -1).mean(dim=1) feats.append(out.cpu().float()) if progress is not None: progress(min(start + img_bs, n) / n) feats = torch.cat(feats, dim=0) if l2_normalize: feats = F.normalize(feats, dim=1) torch.set_rng_state(rng_state) return feats.numpy() # ---------------------------------------------------------------- morphology MORPH_NAMES = [ "R_mean", "R_std", "G_mean", "G_std", "B_mean", "B_std", "dark_frac_0.35", "dark_frac_0.45", "dark_frac_0.55", "grad_mean", "grad_std", "grad_p90", "localvar_mean", "localvar_std", ] def morphology_features(images: List, size: int = 256, progress=None) -> np.ndarray: """Interpretable per-image morphology descriptors for stained micrographs. Captures colour, stain darkness at several thresholds, edge/gradient energy and local-variance granularity -- the properties that separate empty wells, dense colonies, and well-edge tiles. Complements the deep semantic features. """ out = [] n = len(images) for i, x in enumerate(images): im = (Image.open(x) if isinstance(x, str) else x).convert("RGB").resize((size, size)) a = np.asarray(im).astype(np.float32) / 255.0 gray = a.mean(-1) f = [] for c in range(3): f += [float(a[..., c].mean()), float(a[..., c].std())] for thr in (0.35, 0.45, 0.55): f.append(float((gray < thr).mean())) gx = ndimage.sobel(gray, axis=0) gy = ndimage.sobel(gray, axis=1) gmag = np.hypot(gx, gy) f += [float(gmag.mean()), float(gmag.std()), float(np.percentile(gmag, 90))] lv = ndimage.uniform_filter(gray ** 2, 7) - ndimage.uniform_filter(gray, 7) ** 2 f += [float(lv.mean()), float(lv.std())] out.append(f) if progress is not None: progress((i + 1) / n) return np.asarray(out, dtype=np.float32)