inference.py

"""Fusion Perception 1 v0.1 — landmark / place retrieval descriptor.

A frozen DINOv2-L ViT-L/14 backbone (Apache-2.0) plus a small trained projection head
(~3.1M parameters, this repo's only weights). One image goes in, one 512-d L2-normalized
descriptor comes out; rank a gallery by cosine similarity.

    from PIL import Image
    from inference import FusionPerceptionRetrieval

    fp = FusionPerceptionRetrieval.from_pretrained("EximiusLabs/fusion-perception-1-preview")

    gallery = fp.embed([Image.open(p) for p in paths])   # [N, 512]
    query   = fp.embed(Image.open("query.jpg"))          # [512]
    scores, idx = fp.search(query, gallery, topk=10)

The descriptor is multi-scale: the backbone runs at three resolutions (short side 224,
322, 448 px, aspect preserved, long side capped at 1022), the CLS token is L2-normalized
per scale, the three are averaged and re-normalized, and the head maps that 1024-d vector
to 512-d. The backbone is never fine-tuned.

Two heads ship. `standard` is trained on all of GLDv2-clean, the same data published
systems use. `decon` excludes the 87 classes our audit matched to ROxford/RParis query
landmarks. See README_hf.md for what each is for and for the measured difference.

Requires: torch, transformers>=4.46, pillow, numpy, huggingface_hub (only for
from_pretrained on a repo id). The DINOv2 backbone downloads from Meta's repository.

Reported numbers were produced in float16 on CUDA. CPU runs default to float32 and will
differ in the last decimal.
"""
from __future__ import annotations

import json
import os
from typing import Iterable, Optional, Sequence, Tuple, Union

import torch
import torch.nn as nn
import torch.nn.functional as F

DEFAULT_BACKBONE = "facebook/dinov2-large"
SCALES = (1.0, 1.414, 2.0)
BASE_SHORT = 224
MAX_LONG = 1022
PATCH = 14
IMAGENET_MEAN = (0.485, 0.456, 0.406)
IMAGENET_STD = (0.229, 0.224, 0.225)


def _snap(v: float) -> int:
    """Round a side length to a multiple of the patch size (minimum two patches)."""
    return max(PATCH * 2, int(round(v / PATCH)) * PATCH)


def target_size(w: int, h: int, scale: float) -> Tuple[int, int]:
    """Resize geometry for one scale: (width, height), aspect preserved.

    Short side is 224*scale snapped to a multiple of 14; the long side follows the
    aspect ratio, is snapped the same way, and is clamped to 1022.
    """
    short = _snap(BASE_SHORT * scale)
    if w <= h:
        return short, min(_snap(h * short / w), MAX_LONG)
    return min(_snap(w * short / h), MAX_LONG), short


def make_head(emb: int = 512) -> nn.Sequential:
    """The trained head: 1024-d multi-scale CLS -> 512-d descriptor (pre-normalization)."""
    return nn.Sequential(nn.Linear(1024, 2048), nn.GELU(),
                         nn.Linear(2048, emb, bias=False), nn.BatchNorm1d(emb))


class FusionPerceptionRetrieval:
    """Frozen DINOv2-L multi-scale CLS + trained projection head -> 512-d descriptor."""

    def __init__(self, root: str, device: Optional[str] = None,
                 protocol: Optional[str] = None, backbone: Optional[str] = None,
                 dtype: Optional[torch.dtype] = None):
        from transformers import AutoModel

        self.root = root
        with open(os.path.join(root, "config.json")) as fh:
            self.cfg = json.load(fh)
        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
        self.dtype = dtype or (torch.float16 if self.device.startswith("cuda")
                               else torch.float32)
        self.protocol = protocol or self.cfg.get("default_head", "standard")
        heads = self.cfg["heads"]
        if self.protocol not in heads:
            raise ValueError(f"unknown head '{self.protocol}'; available: {list(heads)}")

        self.backbone_id = backbone or self.cfg.get("backbone", DEFAULT_BACKBONE)
        self.model = (AutoModel.from_pretrained(self.backbone_id)
                      .to(device=self.device, dtype=self.dtype).eval())
        for p in self.model.parameters():
            p.requires_grad_(False)

        self.head = make_head(int(self.cfg["descriptor"]["dim"]))
        entry = heads[self.protocol]
        st_name = entry.get("safetensors")
        st_path = os.path.join(root, st_name) if st_name else None
        if st_path and os.path.exists(st_path):
            from safetensors import safe_open

            with safe_open(st_path, framework="pt", device="cpu") as f:
                state = {k: f.get_tensor(k) for k in f.keys()}
                meta = dict(f.metadata() or {})
            self.head.load_state_dict(state)
            # metadata values are strings in safetensors; the JSON-encoded ones decode
            self.head_meta = {}
            for k, v in meta.items():
                try:
                    self.head_meta[k] = json.loads(v)
                except (ValueError, TypeError):
                    self.head_meta[k] = v
        else:
            ck = torch.load(os.path.join(root, entry["file"]),
                            map_location="cpu", weights_only=False)
            self.head.load_state_dict(ck["state_dict"])
            self.head_meta = {k: v for k, v in ck.items() if k != "state_dict"}
        self.head = self.head.to(self.device).float().eval()

        pre = self.cfg.get("preprocess", {})
        self.scales = tuple(pre.get("scales", SCALES))
        mean = torch.tensor(pre.get("mean", IMAGENET_MEAN)).view(1, 3, 1, 1)
        std = torch.tensor(pre.get("std", IMAGENET_STD)).view(1, 3, 1, 1)
        self._mean = mean.to(self.device)
        self._std = std.to(self.device)

    @classmethod
    def from_pretrained(cls, repo_or_path: str, **kw) -> "FusionPerceptionRetrieval":
        if os.path.isdir(repo_or_path):
            root = repo_or_path
        else:
            from huggingface_hub import snapshot_download
            root = snapshot_download(repo_or_path)
        return cls(root, **kw)

    # ------------------------------------------------------------------ backbone
    @torch.no_grad()
    def _cls_at_scale(self, pils: Sequence, scale: float) -> torch.Tensor:
        """L2-normalized CLS token for a batch of PIL images at one scale: [B, 1024].

        Images are grouped by resize shape so a batch of same-shaped images runs in one
        forward pass; mixed aspect ratios simply produce more groups.
        """
        import numpy as np
        from PIL import Image
        groups: dict = {}
        for i, im in enumerate(pils):
            groups.setdefault(target_size(im.size[0], im.size[1], scale), []).append(i)
        out = [None] * len(pils)
        for (nw, nh), idxs in groups.items():
            budget = max(1, 46000 // ((nw // PATCH) * (nh // PATCH) + 5))
            for s0 in range(0, len(idxs), budget):
                grp = idxs[s0:s0 + budget]
                arr = [torch.from_numpy(
                    np.array(pils[i].resize((nw, nh), Image.BICUBIC))
                ).permute(2, 0, 1) for i in grp]
                x = torch.stack(arr).float().to(self.device) / 255.0
                x = ((x - self._mean) / self._std).to(self.dtype)
                hs = self.model(pixel_values=x).last_hidden_state
                c = F.normalize(hs[:, 0].float(), dim=-1).cpu()
                for j, i in enumerate(grp):
                    out[i] = c[j]
        return torch.stack(out)

    @torch.no_grad()
    def embed_backbone(self, images, bbox: Optional[Iterable] = None) -> torch.Tensor:
        """Multi-scale CLS descriptor before the head: [B, 1024], L2-normalized.

        bbox, when given, is one (x1, y1, x2, y2) crop box per image, applied before
        resizing. The revisitop benchmarks crop queries this way; gallery images are
        never cropped.
        """
        one = not isinstance(images, (list, tuple))
        pils = [images] if one else list(images)
        pils = [im.convert("RGB") for im in pils]
        if bbox is not None:
            boxes = [bbox] if one else list(bbox)
            if len(boxes) != len(pils):
                raise ValueError("bbox must have one box per image")
            pils = [im if b is None else im.crop(tuple(int(v) for v in b))
                    for im, b in zip(pils, boxes)]
        acc = torch.stack([self._cls_at_scale(pils, s) for s in self.scales]).mean(0)
        v = F.normalize(acc, dim=-1)
        return v[0] if one else v

    # ------------------------------------------------------------------ descriptor
    @torch.no_grad()
    def embed(self, images, bbox: Optional[Iterable] = None) -> torch.Tensor:
        """512-d L2-normalized retrieval descriptor: [512] for one image, else [B, 512].

        Index these and rank a gallery by cosine similarity (a plain dot product, since
        the vectors are unit norm).
        """
        one = not isinstance(images, (list, tuple))
        ms = self.embed_backbone(images, bbox=bbox)
        if one:
            ms = ms.unsqueeze(0)
        e = self.head(ms.float().to(self.device))
        e = F.normalize(e, dim=-1).cpu()
        return e[0] if one else e

    # ------------------------------------------------------------------ ranking
    @staticmethod
    def search(query: torch.Tensor, gallery: torch.Tensor,
               topk: int = 10) -> Tuple[torch.Tensor, torch.Tensor]:
        """Rank a gallery by cosine similarity. Returns (scores, indices).

        query: [512] or [Q, 512]; gallery: [N, 512]. Both must be L2-normalized, which
        is what embed() returns. topk is clamped to the gallery size.
        """
        q = query.unsqueeze(0) if query.dim() == 1 else query
        sims = q.float() @ gallery.float().T
        scores, idx = torch.topk(sims, k=min(topk, gallery.shape[0]), dim=1)
        if query.dim() == 1:
            return scores[0], idx[0]
        return scores, idx


# Backwards-compatible alias.
FusionPerception = FusionPerceptionRetrieval


if __name__ == "__main__":
    import numpy as np
    from PIL import Image

    root = os.path.dirname(os.path.abspath(__file__))
    fp = FusionPerceptionRetrieval.from_pretrained(root)
    print(f"backbone={fp.backbone_id} head={fp.protocol} "
          f"device={fp.device} dtype={fp.dtype}")

    rng = np.random.default_rng(0)
    imgs = [Image.fromarray((rng.random((h, w, 3)) * 255).astype("uint8"))
            for h, w in ((240, 320), (200, 200), (320, 240))]
    gallery = fp.embed(imgs)
    query = fp.embed(imgs[0])
    scores, idx = fp.search(query, gallery, topk=3)
    print("gallery:", tuple(gallery.shape),
          "| unit norm:", bool(torch.allclose(gallery.norm(dim=-1),
                                              torch.ones(len(imgs)), atol=1e-3)))
    print("query:", tuple(query.shape), "| top-3 idx:", idx.tolist(),
          "| top-3 scores:", [round(s, 4) for s in scores.tolist()])
    assert int(idx[0]) == 0, "an image must retrieve itself first"
    print("OK")