File size: 10,141 Bytes
6fe482b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 | """
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()
|