111 lines
3.4 KiB
Python
111 lines
3.4 KiB
Python
|
from PIL import Image
|
|||
|
|
from transformers import AutoProcessor, AutoModel, SiglipTextConfig
|
|||
|
|
import torch
|
|||
|
|
from typing import List, Union, Optional
|
|||
|
|
import os
|
|||
|
|
|
|||
|
|
class SigLIPInterface:
|
|||
|
|
def __init__(
|
|||
|
|
self,
|
|||
|
|
model_name: str = "google/siglip-base-patch16-224",
|
|||
|
|
max_position_embeddings: int = 64,
|
|||
|
|
device: Optional[str] = None
|
|||
|
|
):
|
|||
|
|
"""
|
|||
|
|
初始化 SigLIP 接口
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
model_name: 模型名称或路径
|
|||
|
|
max_position_embeddings: 最大位置编码长度，默认为64
|
|||
|
|
device: 设备名称 (cuda/cpu)，如果为None则自动选择
|
|||
|
|
"""
|
|||
|
|
self.device = device if device else "cuda" if torch.cuda.is_available() else "cpu"
|
|||
|
|
|
|||
|
|
# 加载模型配置
|
|||
|
|
text_config = SiglipTextConfig.from_pretrained(model_name)
|
|||
|
|
text_config.max_position_embeddings = max_position_embeddings
|
|||
|
|
|
|||
|
|
# 加载模型和处理器
|
|||
|
|
self.model = AutoModel.from_pretrained(
|
|||
|
|
model_name,
|
|||
|
|
text_config=text_config,
|
|||
|
|
ignore_mismatched_sizes=True
|
|||
|
|
).to(self.device)
|
|||
|
|
|
|||
|
|
self.processor = AutoProcessor.from_pretrained(model_name)
|
|||
|
|
|
|||
|
|
def compute_similarity(
|
|||
|
|
self,
|
|||
|
|
image: Union[str, Image.Image],
|
|||
|
|
texts: List[str],
|
|||
|
|
return_probs: bool = True
|
|||
|
|
) -> torch.Tensor:
|
|||
|
|
"""
|
|||
|
|
计算图像和文本的相似度
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
image: 图像路径或 PIL Image 对象
|
|||
|
|
texts: 文本列表
|
|||
|
|
return_probs: 是否返回概率值（True）或原始logits（False）
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
相似度分数张量
|
|||
|
|
"""
|
|||
|
|
# 处理输入图像
|
|||
|
|
if isinstance(image, str):
|
|||
|
|
if not os.path.exists(image):
|
|||
|
|
raise FileNotFoundError(f"Image file not found: {image}")
|
|||
|
|
image = Image.open(image).convert("RGB")
|
|||
|
|
|
|||
|
|
# 处理输入
|
|||
|
|
inputs = self.processor(
|
|||
|
|
text=texts,
|
|||
|
|
images=image,
|
|||
|
|
padding="max_length",
|
|||
|
|
return_tensors="pt"
|
|||
|
|
).to(self.device)
|
|||
|
|
|
|||
|
|
# 计算相似度
|
|||
|
|
with torch.no_grad():
|
|||
|
|
outputs = self.model(**inputs)
|
|||
|
|
logits = outputs.logits_per_image
|
|||
|
|
|
|||
|
|
if return_probs:
|
|||
|
|
return torch.sigmoid(logits)
|
|||
|
|
return logits
|
|||
|
|
|
|||
|
|
def get_similarity_score(
|
|||
|
|
self,
|
|||
|
|
image: Union[str, Image.Image],
|
|||
|
|
text: str
|
|||
|
|
) -> float:
|
|||
|
|
"""
|
|||
|
|
获取单个图像和文本的相似度分数
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
image: 图像路径或 PIL Image 对象
|
|||
|
|
text: 文本字符串
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
相似度分数（0-1之间）
|
|||
|
|
"""
|
|||
|
|
scores = self.compute_similarity(image, [text])
|
|||
|
|
return scores[0][0].item()
|
|||
|
|
|
|||
|
|
# 使用示例
|
|||
|
|
if __name__ == "__main__":
|
|||
|
|
# 创建接口实例
|
|||
|
|
siglip = SigLIPInterface()
|
|||
|
|
|
|||
|
|
# 示例用法
|
|||
|
|
image_path = "image.png"
|
|||
|
|
text = "这里有一个非常精美的盒子。"
|
|||
|
|
|
|||
|
|
# 计算相似度
|
|||
|
|
score = siglip.get_similarity_score(image_path, text)
|
|||
|
|
print(f"相似度分数: {score:.2%}")
|
|||
|
|
|
|||
|
|
# 批量计算
|
|||
|
|
texts = ["这里有一个非常精美的盒子。", "这是一段示例文本", "这里有一个非常精美的商品。"]
|
|||
|
|
scores = siglip.compute_similarity(image_path, texts)
|
|||
|
|
print(f"批量相似度分数: {scores}")
|