bundle pockethb/inference.py

pockethb/inference.py

@@ -0,0 +1,172 @@
+"""Inference pipeline: load HF bundle, embed photos, optionally personalise.
+
+Designed to be the single import-and-use class for chunks 6–8.
+
+    from pockethb.inference import InferenceSession
+    sess = InferenceSession.from_hub()                     # loads bubbaonbubba/pockethb-base
+    raw_hb = sess.predict_aggregate(photo_paths)           # global prediction
+    sess.calibrate(photo_paths, true_hb_g_per_dL=15.3)     # fit affine bias correction
+    personal_hb = sess.predict_aggregate(photo_paths)      # now personalised
+"""
+from __future__ import annotations
+
+import pickle
+from dataclasses import dataclass
+from pathlib import Path
+
+import numpy as np
+import torch
+from PIL import Image
+
+from .calibration import AffineCalibrator, PersonalHead
+from .embed import _prep_crop, load_backbone
+from .preprocess import shades_of_gray
+
+
+def _load_image(src) -> np.ndarray:
+    """Accept str / Path / PIL.Image / numpy array. Return HxWx3 uint8."""
+    if isinstance(src, np.ndarray):
+        return src
+    if isinstance(src, (str, Path)):
+        return np.asarray(Image.open(src).convert("RGB"))
+    if isinstance(src, Image.Image):
+        return np.asarray(src.convert("RGB"))
+    raise TypeError(f"unsupported image source: {type(src)}")
+
+
+@dataclass
+class InferenceResult:
+    raw_per_image: np.ndarray            # global prediction per input photo (g/dL)
+    raw_aggregate: float                 # global prediction at session level (mean+std agg)
+    personal_per_image: np.ndarray | None = None    # post-calibration per photo
+    personal_aggregate: float | None = None         # post-calibration session level
+    method: str = "global"               # "global" | "affine" | "mlp"
+    n_photos: int = 0
+    notes: str = ""
+
+
+class InferenceSession:
+    """Carries the global model bundle + (optional) per-user calibrator."""
+
+    def __init__(self, bundle: dict, device: str = "cpu"):
+        self.bundle = bundle
+        self.backbone_name = bundle["backbone_name"]
+        self.image_size = int(bundle["image_size"])
+        self.sog_p = int(bundle["shades_of_gray_p"])
+        self.blender = bundle["blender"]
+        self.device = device
+        self._backbone = None
+        self.calibrator: AffineCalibrator | None = None
+        self.personal_head: PersonalHead | None = None
+
+    @classmethod
+    def from_hub(cls, repo_id: str = "bubbaonbubba/pockethb-base", device: str = "cpu") -> "InferenceSession":
+        from huggingface_hub import hf_hub_download
+
+        path = hf_hub_download(repo_id=repo_id, filename="pockethb_base.pkl")
+        with open(path, "rb") as f:
+            bundle = pickle.load(f)
+        return cls(bundle, device=device)
+
+    @classmethod
+    def from_pkl(cls, path: str | Path, device: str = "cpu") -> "InferenceSession":
+        with open(path, "rb") as f:
+            bundle = pickle.load(f)
+        return cls(bundle, device=device)
+
+    def _get_backbone(self):
+        if self._backbone is None:
+            self._backbone = load_backbone(self.backbone_name, device=self.device)
+        return self._backbone
+
+    @torch.no_grad()
+    def embed_image(self, image) -> np.ndarray:
+        """Apply Shades-of-Gray + resize + normalise → frozen backbone → 768-d feature."""
+        img = _load_image(image)
+        tensor = _prep_crop(img, apply_sog=True).unsqueeze(0).to(self.device)
+        feat = self._get_backbone()(tensor).cpu().numpy()[0]
+        return feat
+
+    @torch.no_grad()
+    def embed_many(self, images) -> np.ndarray:
+        """Embed a list of images. Returns (n, 768) array."""
+        feats = np.stack([self.embed_image(img) for img in images], axis=0)
+        return feats
+
+    def _aggregate(self, embs: np.ndarray) -> np.ndarray:
+        """Apply the same mean+std per-patient aggregation the global model was trained with."""
+        if embs.ndim == 1:
+            embs = embs[None, :]
+        if embs.shape[0] == 1:
+            agg = np.concatenate([embs[0], np.zeros_like(embs[0])])
+        else:
+            agg = np.concatenate([embs.mean(axis=0), embs.std(axis=0, ddof=0)])
+        return agg.astype(np.float32).reshape(1, -1)
+
+    def predict_per_image(self, images) -> np.ndarray:
+        """Per-image global prediction (each photo treated as its own session)."""
+        embs = self.embed_many(images)
+        preds = []
+        for i in range(embs.shape[0]):
+            agg = self._aggregate(embs[i : i + 1])
+            preds.append(float(self.blender.predict(agg)[0]))
+        return np.array(preds, dtype=np.float64)
+
+    def predict_aggregate(self, images) -> float:
+        """One Hb estimate from a session: aggregate all photos via mean+std and predict once."""
+        embs = self.embed_many(images)
+        agg = self._aggregate(embs)
+        raw = float(self.blender.predict(agg)[0])
+        if self.calibrator and self.calibrator.fitted:
+            return float(self.calibrator.predict(np.array([raw]))[0])
+        return raw
+
+    def calibrate(self, images, true_hb_g_per_dL) -> AffineCalibrator:
+        """Fit per-user affine calibration against a known bloodwork reading.
+
+        true_hb_g_per_dL: scalar (single anchor) or array (multiple paired anchors).
+        """
+        per = self.predict_per_image(images)
+        if np.isscalar(true_hb_g_per_dL):
+            targets = np.full(len(per), float(true_hb_g_per_dL))
+        else:
+            targets = np.asarray(true_hb_g_per_dL, dtype=np.float64).ravel()
+        self.calibrator = AffineCalibrator().fit(per, targets)
+        return self.calibrator
+
+    def calibrate_mlp(self, images, true_hb_g_per_dL, **head_kwargs) -> PersonalHead:
+        """Fit a per-user MLP head on top of the frozen embeddings."""
+        embs = self.embed_many(images)
+        if np.isscalar(true_hb_g_per_dL):
+            targets = np.full(embs.shape[0], float(true_hb_g_per_dL))
+        else:
+            targets = np.asarray(true_hb_g_per_dL, dtype=np.float64).ravel()
+        self.personal_head = PersonalHead(in_dim=embs.shape[1], **head_kwargs).fit(embs, targets)
+        return self.personal_head
+
+    def run(self, images, true_hb_g_per_dL: float | None = None) -> InferenceResult:
+        """Full session-level inference. If true_hb_g_per_dL is given, also fits + applies affine calibration."""
+        raw_per = self.predict_per_image(images)
+        raw_agg = float(np.mean(raw_per))
+
+        if true_hb_g_per_dL is not None:
+            cal = self.calibrate(images, true_hb_g_per_dL)
+            personal_per = cal.predict(raw_per)
+            personal_agg = float(np.mean(personal_per))
+            return InferenceResult(
+                raw_per_image=raw_per,
+                raw_aggregate=raw_agg,
+                personal_per_image=personal_per,
+                personal_aggregate=personal_agg,
+                method=f"affine_{cal.mode}",
+                n_photos=len(raw_per),
+                notes=f"calibrator: a={cal.a:.3f} b={cal.b:+.3f} anchors={cal.n_anchors_used}",
+            )
+
+        return InferenceResult(
+            raw_per_image=raw_per,
+            raw_aggregate=raw_agg,
+            method="global",
+            n_photos=len(raw_per),
+            notes="no calibration applied",
+        )