palimseste/vision.py

"""PALIMPSESTE — Multi-modal image encoder.

The substrate is modality-agnostic: it operates on hypervectors, not on
pixels or tokens. This module bridges the gap by encoding images into the
same bipolar hypervector space ``H = {+1, -1}^D`` that the text LM uses.

How it works
------------
An image is:
  1. resized to a fixed ``grid x grid`` thumbnail (e.g. 16x16 = 256 patches);
  2. each patch is converted to a feature vector (mean RGB, or a flattened
     downscale);
  3. each patch's features are level-coded into hypervectors (using the same
     :class:`Encoder` that text uses);
  4. each patch HV is bound with a **spatial role** (its grid position) and
     bundled — exactly the same sequence/set encoding as text, but with 2D
     roles instead of 1D.

The result is a single D-dimensional hypervector that captures the image's
spatial structure. Two similar images produce similar HVs (quasi-orthogonal
to everything else, correlated to each other), so the associative memory can
retrieve images by content.

This module depends only on numpy (no PIL/torch) — images are passed as
``np.ndarray`` of shape ``(H, W, 3)`` with values in ``[0, 255]``.
"""

from __future__ import annotations

from dataclasses import dataclass, field
import numpy as np

from .hv import HV, bind, bundle, random_hv, bits_to_signs, signs_to_bits, similarity
from .learner import Encoder

__all__ = ["ImageEncoder", "encode_image", "ImageMemory"]


@dataclass
class ImageEncoder:
    """Encodes images into hypervectors using spatial patch binding.

    Parameters
    ----------
    D : int
        Hypervector dimensionality (must match the Memory/Encoder).
    grid : int
        The image is resized to ``grid x grid`` patches.
    encoder : Encoder
        Shared encoder for level-coding patch features (reuses the same
        atom/role tables as text, so text and image HVs live in the same
        space and can cross-reference).
    """

    D: int
    grid: int = 16
    encoder: Encoder = field(init=False)

    def __post_init__(self) -> None:
        self.encoder = Encoder(D=self.D)

    def _spatial_role(self, row: int, col: int) -> HV:
        """A deterministic spatial role HV for grid position (row, col)."""
        return self.encoder._atom(("spatial", row, col))

    def _patch_features(self, patch: np.ndarray) -> np.ndarray:
        """Extract a feature vector from a patch (mean RGB + brightness)."""
        # patch shape: (ph, pw, 3)
        mean_rgb = patch.reshape(-1, 3).mean(axis=0) / 255.0  # (3,) in [0,1]
        brightness = float(mean_rgb.mean())
        return np.append(mean_rgb, brightness)  # (4,)

    def _features_to_hv(self, features: np.ndarray) -> HV:
        """Level-code each feature and bundle them into one patch HV."""
        hvs = []
        for i, val in enumerate(features):
            # use encode_float with role binding for each feature channel
            fhv = self.encoder.encode_float(float(val), lo=0.0, hi=1.0)
            role = self.encoder._role(i)
            hvs.append(bind(fhv, role))
        return bundle(hvs, rng=self.encoder.rng, deterministic=True)

    def encode(self, image: np.ndarray) -> HV:
        """Encode an image ``(H, W, 3)`` uint8 into a hypervector.

        The image is downsampled to ``grid x grid`` patches, each patch is
        encoded into a feature HV, and patches are bound with their spatial
        roles and bundled into a single image HV.
        """
        if image.ndim != 3 or image.shape[2] != 3:
            raise ValueError(f"image must be (H, W, 3), got {image.shape}")

        H, W = image.shape[:2]
        gh, gw = H // self.grid, W // self.grid
        if gh == 0 or gw == 0:
            # image too small; upscale via repeat
            image = np.repeat(np.repeat(image, self.grid // max(H, 1) + 1, axis=0),
                              self.grid // max(W, 1) + 1, axis=1)
            H, W = image.shape[:2]
            gh, gw = H // self.grid, W // self.grid

        patch_hvs: list[HV] = []
        for r in range(self.grid):
            for c in range(self.grid):
                patch = image[r * gh:(r + 1) * gh, c * gw:(c + 1) * gw]
                if patch.size == 0:
                    continue
                feats = self._patch_features(patch)
                patch_hv = self._features_to_hv(feats)
                role = self._spatial_role(r, c)
                patch_hvs.append(bind(patch_hv, role))

        if not patch_hvs:
            return random_hv(self.D, rng=self.encoder.rng)
        return bundle(patch_hvs, rng=self.encoder.rng, deterministic=True)


def encode_image(image: np.ndarray, D: int = 10000, grid: int = 16) -> HV:
    """Convenience: encode an image into a hypervector."""
    enc = ImageEncoder(D=D, grid=grid)
    return enc.encode(image)


@dataclass
class ImageMemory:
    """An associative image memory: store and retrieve images by HV similarity.

    Uses a :class:`Memory` under the hood, so images are append-only and
    never forgotten. Retrieve by query HV (from another image, or from a
    text description if cross-modal roles are shared).
    """

    D: int
    encoder: ImageEncoder = field(init=False)
    _images: list[np.ndarray] = field(default_factory=list)
    _labels: list[str] = field(default_factory=list)
    _hvs: list[HV] = field(default_factory=list)

    def __post_init__(self) -> None:
        self.encoder = ImageEncoder(D=self.D)

    def store(self, image: np.ndarray, label: str = "") -> HV:
        """Store an image, return its hypervector."""
        hv = self.encoder.encode(image)
        self._images.append(image)
        self._labels.append(label)
        self._hvs.append(hv)
        return hv

    def retrieve(self, query_hv: HV, top_k: int = 1) -> list[tuple[int, float, str]]:
        """Find the ``top_k`` most similar stored images.

        Returns list of ``(index, similarity, label)`` sorted by similarity.
        """
        if not self._hvs:
            return []
        sims = [(i, similarity(query_hv, h), self._labels[i])
                for i, h in enumerate(self._hvs)]
        sims.sort(key=lambda x: -x[1])
        return sims[:top_k]

    def query_by_image(self, image: np.ndarray, top_k: int = 1):
        """Find similar stored images."""
        qhv = self.encoder.encode(image)
        return self.retrieve(qhv, top_k=top_k)

    @property
    def size(self) -> int:
        return len(self._images)