inference.py

"""Tactus Mat inference: open-vocabulary posture recognition from a body pressure mat.

Tactus Mat embeds a window of 64x32 pressure frames into the fusion-embedding canonical
2048-d text space, so recognition is a cosine ranking against text queries rather than a
fixed classifier head. It is the bed/seat/body-mat sensor profile of the Tactus tactile
pack; the glove profile lives at EximiusLabs/fusion-embedding-2-tactus.

Two halves, loaded from two places:
  * PRESSURE side (this repository): a trained CNN trunk + projector
    (``tactus_mat_head.pt`` / ``model.safetensors``), defined in the vendored ``tactile.py``
    (byte-identical to the training-time module, so the checkpoint loads with strict=True).
  * TEXT side (the base model repository): the canonical whitened text readout of
    ``EximiusLabs/fusion-embedding-2-2b-preview`` via the ``fusion-embedding`` package.
    The model was trained against that exact readout; embedding text any other way
    (including the plain base model without whitening) will misrank.

Input contract: frames are pressure maps scaled to [0,1]. The training data is a Vista
Medical FSA SoftFlex 2048 (32x64 sensels) whose documented range is [0, 1000]; convert raw
counts with ``mat_raw_to_unit`` or pass ``raw="fsa"``. A window is any number of frames;
8 evenly spaced frames of a held posture is what training used.

Usage:
    from inference import TactusMatEmbedder

    tm = TactusMatEmbedder.from_pretrained("EximiusLabs/fusion-embedding-2-tactus-mat",
                                           revision="v0.1-preview")
    window = np.load("mat_window.npy")       # [F,64,32] or a single [64,32] frame
    for text, score in tm.rank(window, ["a person lying on their left side",
                                        "a person lying flat on their back",
                                        "a person in the fetal position"]):
        print(f"{score:+.3f}  {text}")
"""
from __future__ import annotations

import os

import numpy as np
import torch
import torch.nn.functional as F

from tactile import TactileEncoder, build_projector       # vendored, verbatim

REPO_DEFAULT = "EximiusLabs/fusion-embedding-2-tactus-mat"
BASE_REPO = "EximiusLabs/fusion-embedding-2-2b-preview"
SAFETENSORS_FILE = "model.safetensors"
CKPT_FILE = "tactus_mat_head.pt"

# Vista Medical FSA SoftFlex 2048 documented sensor range (PhysioNet pmd/1.0.0).
FSA_RANGE_MAX = 1000.0
GRID = (64, 32)


def mat_raw_to_unit(raw) -> np.ndarray:
    """Raw mat counts -> [0,1] floats (clip(raw/1000, 0, 1), the training normalization)."""
    a = np.asarray(raw, dtype=np.float32)
    return np.clip(a / FSA_RANGE_MAX, 0.0, 1.0)


def preprocess_mat(frames) -> torch.Tensor:
    """Accept [64,32] or [F,64,32] (numpy/list/tensor) -> float32 [F,64,32] in [0,1].

    uint8 0-255 input is scaled by 1/255. Float input already in [0,1] passes through;
    float input exceeding 1.0 is treated as raw counts and rescaled by the FSA range.
    """
    if isinstance(frames, torch.Tensor):
        t = frames.detach().float()
    else:
        t = torch.as_tensor(np.ascontiguousarray(np.asarray(frames)), dtype=torch.float32)
    if t.dim() == 2:
        t = t.unsqueeze(0)
    if t.dim() != 3:
        raise ValueError(f"expected [64,32] or [F,64,32], got {list(t.shape)}")
    if tuple(t.shape[-2:]) != GRID:
        raise ValueError(f"expected a {GRID[0]}x{GRID[1]} pressure grid "
                         f"(rows x cols), got {list(t.shape[-2:])}")
    mx = float(t.max()) if t.numel() else 0.0
    if mx > 1.0:
        t = t / (255.0 if mx <= 255.0 else FSA_RANGE_MAX)
    return t.clamp_(0.0, 1.0)


class TactusMatEmbedder:
    """Mat pressure -> 2048-d canonical embeddings, plus the matching text side."""

    def __init__(self, head_path: str, device: str = "cuda", dtype=torch.bfloat16,
                 base_repo: str = BASE_REPO, load_text: bool = True):
        self.device = device
        if head_path.endswith(".safetensors"):
            import json as _json

            import safetensors.torch as st
            from safetensors import safe_open
            flat = st.load_file(head_path)
            with safe_open(head_path, framework="pt") as f:
                cfg = _json.loads((f.metadata() or {})["config"])
            enc_sd = {k[len("encoder."):]: v for k, v in flat.items()
                      if k.startswith("encoder.")}
            proj_sd = {k[len("proj."):]: v for k, v in flat.items() if k.startswith("proj.")}
        else:
            blob = torch.load(head_path, map_location="cpu", weights_only=False)
            cfg = blob["config"]
            enc_sd = {k: v.float() for k, v in blob["encoder"].items()}
            proj_sd = {k: v.float() for k, v in blob["proj"].items()}
        self.cfg = cfg
        self.window_frames = int(cfg.get("K", cfg.get("window_frames", 8)))

        self.enc = TactileEncoder(temporal=cfg.get("temporal", "conv"),
                                  depth=cfg.get("depth", "resnet18"),
                                  frames=self.window_frames,
                                  flow=bool(cfg.get("flow", False)))
        self.proj = build_projector()
        self.enc.load_state_dict(enc_sd, strict=True)
        self.proj.load_state_dict(proj_sd, strict=True)
        self.enc.eval().to(device)
        self.proj.eval().to(device)

        self.ue = None
        if load_text:
            from fusion_embedding import UnifiedEmbedder
            self.ue = UnifiedEmbedder.from_pretrained(base_repo, device=device, dtype=dtype)

    @classmethod
    def from_pretrained(cls, repo: str = REPO_DEFAULT, revision: str | None = None,
                        device: str = "cuda", dtype=torch.bfloat16, load_text: bool = True):
        if os.path.isdir(repo):
            head = None
            for name in (SAFETENSORS_FILE, CKPT_FILE):
                for cand in (os.path.join(repo, name), os.path.join(repo, "out", name)):
                    if os.path.exists(cand):
                        head = cand
                        break
                if head:
                    break
            assert head, f"no {SAFETENSORS_FILE} or {CKPT_FILE} under {repo}"
        else:
            from huggingface_hub import hf_hub_download
            try:
                head = hf_hub_download(repo, SAFETENSORS_FILE, revision=revision)
            except Exception:  # noqa: BLE001
                head = hf_hub_download(repo, CKPT_FILE, revision=revision)
        return cls(head, device=device, dtype=dtype, load_text=load_text)

    @torch.no_grad()
    def embed_pressure(self, frames, raw: str | None = None) -> torch.Tensor:
        """A window ``[F,64,32]`` (or one ``[64,32]`` frame) -> 2048-d, L2-normalized."""
        if raw == "fsa":
            frames = mat_raw_to_unit(frames)
        elif raw is not None:
            raise ValueError(f"unknown raw mode {raw!r}; use raw='fsa' or pre-normalize")
        x = preprocess_mat(frames).to(self.device)             # [F,64,32] in [0,1]
        feat = self.enc(x)                                     # [512]
        vec = self.proj(feat)                                  # [2048] raw
        return F.normalize(vec.float(), dim=-1).cpu()

    @torch.no_grad()
    def embed_text(self, texts) -> torch.Tensor:
        """Text -> canonical 2048-d embeddings (whitened readout). Requires load_text=True."""
        assert self.ue is not None, "constructed with load_text=False"
        if isinstance(texts, str):
            texts = [texts]
        return torch.stack([self.ue.embed_text(t).float() for t in texts])

    @torch.no_grad()
    def rank(self, frames, texts, raw: str | None = None):
        """Rank candidate texts against one pressure window; returns [(text, cosine)] sorted."""
        p = self.embed_pressure(frames, raw=raw)
        t = F.normalize(self.embed_text(texts), dim=-1)
        sims = (t @ p).tolist()
        return sorted(zip(list(texts), sims), key=lambda kv: -kv[1])