""" FLAIR foundation-model integration. FLAIR (Silva-Rodriguez et al., Medical Image Analysis 2024) is a CLIP-style vision-language model of the retina: a ResNet-50 image encoder aligned with a Bio-ClinicalBERT text encoder. We use it for: * zero-shot disease classification (image-text similarity over a disease set), * a lightweight VQA / open-vocabulary probe (score arbitrary clinical phrases), * FLAIR embeddings as the feature space for batch-effect analysis, * a quality-vs-disease *disentanglement* readout - measuring whether FLAIR's disease confidence is being contaminated by image degradation (the "does the foundation model conflate degradation with pathology?" question). The model weights are downloaded on first use from the FLAIR Hugging Face repo. Loading is lazy and defensive: if torch / weights are unavailable the rest of the QC tool keeps working, and the FLAIR tabs show a clear status message. """ from __future__ import annotations import numpy as np import threading # Disease vocabulary for zero-shot (a clinically meaningful default set drawn # from FLAIR's own category dictionary). DEFAULT_DISEASES = [ "normal", "diabetic retinopathy", "diabetic macular edema", "age related macular degeneration", "glaucoma", "media haze", "drusens", "pathologic myopia", "branch retinal vein occlusion", "tessellation", "epiretinal membrane", "macular hole", "optic disc cupping", ] # Quality-probe phrases: FLAIR similarity to these tells us whether the *image* # reads as degraded independent of disease. QUALITY_PHRASES_GOOD = ["a clear fundus photograph", "good quality retinal image", "sharp well focused fundus"] QUALITY_PHRASES_BAD = ["a blurred fundus photograph", "noisy low quality image", "ungradable fundus image", "no fundus"] class FlairEngine: """Thread-safe lazy wrapper around FLAIRModel.""" def __init__(self, repo="jusiro2/FLAIR", local_weights=None): self.repo = repo self.local_weights = local_weights self.model = None self.status = "not loaded" self.available = None self._lock = threading.Lock() # ------------------------------------------------------------------ loading def load(self): with self._lock: if self.model is not None: return True try: import torch # noqa: F401 from flair import FLAIRModel if self.local_weights: self.model = FLAIRModel(from_checkpoint=True, weights_path=self.local_weights) else: self.model = FLAIRModel.from_pretrained(self.repo) self.model.eval() self.available = True self.status = "loaded" return True except Exception as e: # torch missing, no weights, offline, etc. self.available = False self.status = f"unavailable ({type(e).__name__}: {e})" return False def is_ready(self): return self.model is not None # -------------------------------------------------------------- inference def _forward(self, rgb, phrases): """Return (probs, logits) for one image over a phrase list.""" img = np.array(rgb) probs, logits = self.model(img, list(phrases)) return np.asarray(probs).ravel(), np.asarray(logits).ravel() def zero_shot_disease(self, rgb, diseases=None, domain_knowledge=True): """Zero-shot disease probabilities for one image.""" if not self.load(): return None diseases = diseases or DEFAULT_DISEASES probs, logits = self._forward(rgb, diseases) order = np.argsort(-probs) return [dict(label=diseases[i], prob=float(probs[i]), logit=float(logits[i])) for i in order] def phrase_similarity(self, rgb, phrases): """Raw open-vocabulary image-text similarity for arbitrary phrases.""" if not self.load(): return None probs, logits = self._forward(rgb, phrases) order = np.argsort(-logits) return [dict(phrase=phrases[i], logit=float(logits[i]), prob=float(probs[i])) for i in order] def ungradable_prob(self, rgb): """Probability the image reads as ungradable/low-quality to FLAIR.""" if not self.load(): return float("nan") probs, _ = self._forward(rgb, ["a clear gradable fundus photograph", "an ungradable low quality fundus image"]) return float(probs[1]) # ------------------------------------------------------------ true VQA # FLAIR is contrastive (not generative), so VQA = pick the best answer among # candidate hypotheses. We turn a question into an answer set (either # user-supplied, or via templates) and rank the answers by image-text match. QUESTION_TEMPLATES = { "gradable": (["a gradable fundus image", "an ungradable fundus image"], ["is this gradable", "can this be graded", "quality ok"]), "disease": (None, ["what disease", "what condition", "diagnosis", "what is wrong", "pathology", "what does this show", "what can you see"]), "dr": (["no diabetic retinopathy", "mild diabetic retinopathy", "moderate diabetic retinopathy", "severe diabetic retinopathy", "proliferative diabetic retinopathy"], ["diabetic retinopathy grade", "dr severity", "how severe"]), "blur": (["a sharp in-focus fundus image", "a blurred out-of-focus image"], ["is it blurred", "in focus", "sharp or blurry"]), "laterality": (["a right eye OD fundus", "a left eye OS fundus"], ["which eye", "left or right", "laterality", "od or os"]), } def vqa_answer(self, rgb, question, candidates=None): """Answer a free-text question by ranking candidate answers. If `candidates` is given (comma/newline separated), those are the answer set; otherwise the question is matched to a template answer set, falling back to the disease vocabulary.""" if not self.load(): return None q = (question or "").lower().strip() answers, source = None, "custom" if candidates: answers = [c.strip() for c in candidates.replace("\n", ",").split(",") if c.strip()] if not answers: for key, (ans, keys) in self.QUESTION_TEMPLATES.items(): if any(k in q for k in keys): if ans is None: # disease-type question answers = list(DEFAULT_DISEASES); source = "disease vocabulary" else: answers = ans; source = f"{key} template" break if not answers: answers = list(DEFAULT_DISEASES); source = "disease vocabulary (default)" probs, logits = self._forward(rgb, answers) order = np.argsort(-probs) ranked = [dict(answer=answers[i], prob=float(probs[i]), logit=float(logits[i])) for i in order] return dict(answer=ranked[0]["answer"], confidence=ranked[0]["prob"], ranked=ranked, answer_set=source) def embedding(self, rgb): """L2-normalised FLAIR image embedding (projection space) for one image.""" if not self.load(): return None import torch with torch.no_grad(): x = self.model.preprocess_image(np.array(rgb)) emb = self.model.vision_model(x) return emb.cpu().numpy().ravel() # ------------------------------------------------ quality/disease disentangle def quality_disentanglement(self, rgb, qc_composite): """Compare FLAIR's *image-quality* read to FLAIR's *disease* confidence, and to the geometric QC composite, to expose whether disease confidence is being driven by degradation rather than pathology. Returns a dict with: flair_quality : 0-100, FLAIR's own good-vs-bad-quality readout flair_disease : top non-normal disease + its probability contamination : signed indicator; high positive == FLAIR calls disease while both QC and FLAIR-quality say the image is degraded (suggesting confound, not true pathology). """ if not self.load(): return None good_l = self._forward(rgb, QUALITY_PHRASES_GOOD)[1] bad_l = self._forward(rgb, QUALITY_PHRASES_BAD)[1] # softmax over [mean good, mean bad] g, b = good_l.mean(), bad_l.mean() flair_quality = float(np.exp(g) / (np.exp(g) + np.exp(b)) * 100) dz = self.zero_shot_disease(rgb) non_normal = [d for d in dz if d["label"] != "normal"] top = non_normal[0] if non_normal else dict(label="normal", prob=0.0) # Contamination heuristic: FLAIR confident in disease AND image is # degraded on both independent quality readouts. degraded = (100 - qc_composite) / 100 # geometric QC says bad flair_bad = (100 - flair_quality) / 100 # FLAIR says bad quality contamination = float(top["prob"] * 0.5 * (degraded + flair_bad)) return dict( flair_quality=flair_quality, flair_disease=top["label"], flair_disease_prob=float(top["prob"]), qc_composite=float(qc_composite), contamination=contamination, note=("High: disease confidence may be driven by degradation, not " "pathology - interpret with caution." if contamination > 0.25 else "Low: disease read appears independent of image quality."), ) # Module-level singleton so the model is loaded once per process. ENGINE = FlairEngine()