src/flair_wrapper.py

"""
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()