pockethb/inference.py

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