File size: 6,620 Bytes
cf29823 f2487c7 cf29823 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 | """Pretrained vision-model embeddings for small-seal glyph retrieval."""
from __future__ import annotations
from pathlib import Path
from typing import Iterable
import numpy as np
import torch
import torch.nn as nn
from PIL import Image
from transformers import AutoImageProcessor, AutoModel
from torchvision.models import resnet50
DEFAULT_VISUAL_MODEL = "facebook/dinov2-small"
def load_visual_model(model_id: str, device: str = "cpu"):
"""Load a Hugging Face vision backbone in inference-only mode."""
processor = AutoImageProcessor.from_pretrained(model_id)
model = AutoModel.from_pretrained(model_id)
model.to(device)
model.eval()
return processor, model
def embed_images(
image_paths: Iterable[Path], processor, model, batch_size: int = 16, device: str = "cpu"
) -> np.ndarray:
"""Create L2-normalized CLS embeddings for a sequence of glyph images."""
paths = list(image_paths)
chunks: list[np.ndarray] = []
with torch.inference_mode():
for start in range(0, len(paths), batch_size):
batch_paths = paths[start : start + batch_size]
images = []
for path in batch_paths:
with Image.open(path) as image:
images.append(image.convert("RGB"))
inputs = processor(images=images, return_tensors="pt")
inputs = {key: value.to(device) for key, value in inputs.items()}
outputs = model(**inputs)
embedding = getattr(outputs, "pooler_output", None)
if embedding is None:
embedding = outputs.last_hidden_state[:, 0]
embedding = torch.nn.functional.normalize(embedding, dim=1)
chunks.append(embedding.cpu().numpy().astype(np.float32))
return np.concatenate(chunks, axis=0)
def load_local_resnet50_feature_model(
checkpoint_path: Path, num_classes: int, device: str = "cpu"
) -> nn.Module:
"""Load this project's classifier and expose its 2048-D pre-classifier features."""
model = resnet50(weights=None)
model.fc = nn.Sequential(nn.Dropout(0.3), nn.Linear(model.fc.in_features, num_classes))
try:
checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=True)
except Exception:
checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=False)
if isinstance(checkpoint, nn.Module):
state_dict = checkpoint.state_dict()
elif isinstance(checkpoint, dict):
state_dict = checkpoint.get("state_dict", checkpoint.get("model_state", checkpoint))
else:
raise ValueError("不支持的模型文件格式")
model.load_state_dict(state_dict, strict=True)
feature_model = nn.Sequential(*list(model.children())[:-1]).to(device)
feature_model.eval()
return feature_model
def load_seal_retriever_feature_model(checkpoint_path: Path, device: str = "cpu") -> nn.Module:
"""Load a small-seal fine-tuned backbone and expose 2048-D features.
The training checkpoint intentionally stores no classification head, so it
can be used for retrieval across both supervised and single-sample chars.
"""
model = resnet50(weights=None)
try:
checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=True)
except Exception:
checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=False)
state_dict = checkpoint.get("backbone_state_dict", checkpoint) if isinstance(checkpoint, dict) else checkpoint
if not isinstance(state_dict, dict):
raise ValueError("不支持的小篆检索训练权重格式")
incompatible = model.load_state_dict(state_dict, strict=False)
unexpected = [key for key in incompatible.unexpected_keys if not key.startswith("fc.")]
if unexpected:
raise ValueError(f"小篆检索权重包含未知参数：{unexpected[0]}")
feature_model = nn.Sequential(*list(model.children())[:-1]).to(device)
feature_model.eval()
return feature_model
def embed_images_local_resnet50(
image_paths: Iterable[Path],
model: nn.Module,
batch_size: int = 16,
device: str = "cpu",
image_size: int = 128,
) -> np.ndarray:
"""Create L2-normalized embeddings with the project's existing classifier backbone."""
paths = list(image_paths)
chunks: list[np.ndarray] = []
mean = torch.tensor([0.485, 0.456, 0.406], dtype=torch.float32).view(3, 1, 1)
std = torch.tensor([0.229, 0.224, 0.225], dtype=torch.float32).view(3, 1, 1)
with torch.inference_mode():
for start in range(0, len(paths), batch_size):
images = []
for path in paths[start : start + batch_size]:
with Image.open(path) as image:
rgb = image.convert("RGB").resize(
(image_size, image_size), Image.Resampling.BILINEAR
)
pixels = np.asarray(rgb, dtype=np.float32) / 255.0
tensor = torch.from_numpy(pixels).permute(2, 0, 1)
images.append((tensor - mean) / std)
batch = torch.stack(images).to(device)
embedding = model(batch).flatten(1)
embedding = torch.nn.functional.normalize(embedding, dim=1)
chunks.append(embedding.cpu().numpy().astype(np.float32))
return np.concatenate(chunks, axis=0)
def embed_pil_images_local_resnet50(
images: Iterable[Image.Image],
model: nn.Module,
batch_size: int = 16,
device: str = "cpu",
image_size: int = 128,
) -> np.ndarray:
"""Embed in-memory images using the same preprocessing as the index builder."""
image_list = list(images)
chunks: list[np.ndarray] = []
mean = torch.tensor([0.485, 0.456, 0.406], dtype=torch.float32).view(3, 1, 1)
std = torch.tensor([0.229, 0.224, 0.225], dtype=torch.float32).view(3, 1, 1)
with torch.inference_mode():
for start in range(0, len(image_list), batch_size):
tensors = []
for image in image_list[start : start + batch_size]:
rgb = image.convert("RGB").resize(
(image_size, image_size), Image.Resampling.BILINEAR
)
pixels = np.asarray(rgb, dtype=np.float32) / 255.0
tensor = torch.from_numpy(pixels).permute(2, 0, 1)
tensors.append((tensor - mean) / std)
batch = torch.stack(tensors).to(device)
embedding = torch.nn.functional.normalize(model(batch).flatten(1), dim=1)
chunks.append(embedding.cpu().numpy().astype(np.float32))
return np.concatenate(chunks, axis=0)
|