src/mcss/vision_probe/features.py
"""Frozen DINOv2 features and patch-level RGB observations for the 2D probe."""
from __future__ import annotations
from collections.abc import Mapping
from pathlib import Path
from typing import TypedDict
import numpy as np
import torch
from PIL import Image
from torch import Tensor
class FeatureResult(TypedDict):
"""CPU tensors and image metadata produced for one image."""
features: Tensor
rgb: Tensor
original_size: tuple[int, int]
image_size: tuple[int, int]
class FrozenDinoExtractor:
"""Load a locally pinned DINOv2-S/14 checkpoint and extract patch features."""
patch_size = 14
feature_dim = 384
def __init__(
self,
repo_root: Path,
weights: Path,
device: str | torch.device = "cuda",
size: int = 224,
) -> None:
self.repo_root = _existing_directory(repo_root, "DINOv2 repository")
self.weights = _existing_file(weights, "DINOv2 checkpoint")
self.device = _validate_device(device)
if isinstance(size, bool) or not isinstance(size, int) or size <= 0:
raise ValueError("size must be a positive integer")
if size % self.patch_size != 0:
raise ValueError(f"size must be divisible by the {self.patch_size}-pixel patch size")
self.size = size
# pretrained=False is intentional: all weights must come from the explicit checkpoint.
model = torch.hub.load(
str(self.repo_root),
"dinov2_vits14",
source="local",
pretrained=False,
)
state_dict = torch.load(self.weights, map_location="cpu", weights_only=True)
if not isinstance(state_dict, Mapping):
raise ValueError("DINOv2 checkpoint must be a raw state_dict mapping")
model.load_state_dict(state_dict, strict=True)
model_dim = getattr(model, "embed_dim", None)
if model_dim is not None and int(model_dim) != self.feature_dim:
raise ValueError(
f"dinov2_vits14 must have feature dimension {self.feature_dim}, found {model_dim}"
)
model_patch_size = getattr(model, "patch_size", self.patch_size)
if isinstance(model_patch_size, tuple):
if model_patch_size != (self.patch_size, self.patch_size):
raise ValueError("dinov2_vits14 must use 14x14 patches")
elif model_patch_size != self.patch_size:
raise ValueError("dinov2_vits14 must use 14x14 patches")
self.model = model.to(self.device).eval()
for parameter in self.model.parameters():
parameter.requires_grad_(False)
def extract(self, path: Path) -> FeatureResult:
"""Extract frozen patch descriptors and the mean RGB of each image patch."""
image_path = _existing_file(path, "image")
with Image.open(image_path) as image:
image = image.convert("RGB")
original_size = (int(image.width), int(image.height))
resized = image.resize((self.size, self.size), resample=Image.Resampling.BICUBIC)
rgb_image = torch.from_numpy(np.asarray(resized, dtype=np.float32)).div_(255.0)
if rgb_image.shape != (self.size, self.size, 3):
raise ValueError(f"resized RGB image has unexpected shape {tuple(rgb_image.shape)}")
rgb = _patch_mean_rgb(rgb_image, self.patch_size)
image_tensor = rgb_image.permute(2, 0, 1).unsqueeze(0)
mean = torch.tensor((0.485, 0.456, 0.406), dtype=torch.float32).view(1, 3, 1, 1)
std = torch.tensor((0.229, 0.224, 0.225), dtype=torch.float32).view(1, 3, 1, 1)
image_tensor = ((image_tensor - mean) / std).to(self.device)
with torch.inference_mode():
output = self.model.forward_features(image_tensor)
if not isinstance(output, Mapping) or "x_norm_patchtokens" not in output:
raise ValueError("DINOv2 forward_features must return x_norm_patchtokens")
patch_tokens = output["x_norm_patchtokens"]
if not isinstance(patch_tokens, Tensor):
raise ValueError("DINOv2 x_norm_patchtokens must be a tensor")
grid = self.size // self.patch_size
expected_tokens = grid * grid
if patch_tokens.ndim != 3 or patch_tokens.shape[0] != 1:
raise ValueError(
"DINOv2 x_norm_patchtokens must have shape [1, tokens, feature_dim]"
)
if patch_tokens.shape[1] != expected_tokens:
raise ValueError(
f"expected {expected_tokens} patch tokens for size {self.size}, "
f"found {patch_tokens.shape[1]}"
)
if patch_tokens.shape[2] != self.feature_dim:
raise ValueError(
f"expected feature dimension {self.feature_dim}, found {patch_tokens.shape[2]}"
)
if not torch.isfinite(patch_tokens).all():
raise ValueError("DINOv2 patch features must contain only finite values")
features = (
patch_tokens[0]
.reshape(grid, grid, self.feature_dim)
.detach()
.to(device="cpu", dtype=torch.float32)
.contiguous()
)
return {
"features": features,
"rgb": rgb.to(device="cpu", dtype=torch.float32).contiguous(),
"original_size": original_size,
"image_size": (self.size, self.size),
}
def _existing_directory(value: Path, label: str) -> Path:
try:
path = Path(value).expanduser().resolve()
except (TypeError, ValueError) as error:
raise TypeError(f"{label} must be a filesystem path") from error
if not path.is_dir():
raise FileNotFoundError(f"{label} does not exist or is not a directory: {path}")
return path
def _existing_file(value: Path, label: str) -> Path:
try:
path = Path(value).expanduser().resolve()
except (TypeError, ValueError) as error:
raise TypeError(f"{label} must be a filesystem path") from error
if not path.is_file():
raise FileNotFoundError(f"{label} does not exist or is not a file: {path}")
return path
def _validate_device(value: str | torch.device) -> torch.device:
try:
device = torch.device(value)
except (TypeError, RuntimeError) as error:
raise ValueError(f"invalid torch device: {value!r}") from error
if device.type not in {"cpu", "cuda"}:
raise ValueError("device must be cpu or cuda")
if device.type == "cuda":
if not torch.cuda.is_available():
raise RuntimeError("CUDA device requested but CUDA is unavailable")
if device.index is not None and device.index >= torch.cuda.device_count():
raise ValueError(f"CUDA device index {device.index} is unavailable")
return device
def _patch_mean_rgb(rgb_image: Tensor, patch_size: int) -> Tensor:
"""Average RGB over non-overlapping image patches, preserving H-W-C order."""
height, width, channels = rgb_image.shape
if channels != 3 or height % patch_size != 0 or width % patch_size != 0:
raise ValueError("RGB image dimensions must be divisible by the patch size")
patches = rgb_image.unfold(0, patch_size, patch_size).unfold(1, patch_size, patch_size)
return patches.mean(dim=(-1, -2))