📚 课程目录第4章 BLIP视觉编码器特征提取

4. BLIP视觉编码器特征提取

任务说明

本案例使用BLIP模型的视觉编码器（ViT-B/16）对一组图像提取全局特征（[CLS] 对应的输出向量），并基于余弦相似度实现图像检索。输入为本地图像目录，输出为查询图像的Top-K相似图像列表及相似度分数。

实现方案

利用HuggingFace transformers库中的BlipModel，该模型包含独立的视觉编码器。通过get_image_features方法直接获取归一化后的全局特征向量，无需额外池化层。检索时计算查询特征与库特征的余弦相似度（特征已归一化，点积即余弦相似度），按降序返回Top-K结果。

代码实现

import torch
import torch.nn.functional as F
from PIL import Image
from transformers import BlipModel, BlipProcessor
import os
import glob

# ========== 配置 ==========
MODEL_NAME = "Salesforce/blip-image-captioning-base"
IMAGE_DIR = "./images"          # 图像库目录（支持jpg/png）
QUERY_IMAGE = "./query.jpg"     # 查询图像
TOP_K = 3                       # 返回最相似的K张

# ========== 初始化模型 ==========
device = "cuda" if torch.cuda.is_available() else "cpu"
processor = BlipProcessor.from_pretrained(MODEL_NAME)
model = BlipModel.from_pretrained(MODEL_NAME).to(device).eval()

# ========== 提取单张图像特征 ==========
def extract_feature(image_path):
    image = Image.open(image_path).convert("RGB")
    inputs = processor(images=image, return_tensors="pt").to(device)
    with torch.no_grad():
        # 返回归一化的图像特征 [1, 256]
        feature = model.get_image_features(**inputs)
    return feature.cpu().squeeze(0)  # [256]

# ========== 构建图像库特征 ==========
image_paths = sorted(glob.glob(os.path.join(IMAGE_DIR, "*.jpg")) +
                     glob.glob(os.path.join(IMAGE_DIR, "*.png")))
assert len(image_paths) > 0, f"图像目录 {IMAGE_DIR} 中没有jpg/png文件"

gallery_features = []
for path in image_paths:
    feat = extract_feature(path)
    gallery_features.append(feat)
gallery_features = torch.stack(gallery_features)  # [N, 256]

# ========== 图像检索 ==========
query_feat = extract_feature(QUERY_IMAGE)
# 特征已归一化，点积即余弦相似度
similarities = torch.mv(gallery_features, query_feat)  # [N]
topk_scores, topk_indices = torch.topk(similarities, k=min(TOP_K, len(image_paths)))

# ========== 输出结果 ==========
print(f"查询图像: {QUERY_IMAGE}")
print(f"图像库数量: {len(image_paths)}")
print(f"Top-{len(topk_indices)} 检索结果:")
for rank, (idx, score) in enumerate(zip(topk_indices, topk_scores), 1):
    print(f"  {rank}. {image_paths[idx]}  相似度: {score.item():.4f}")

运行与验证

依赖安装：

pip install torch torchvision transformers pillow

准备数据：创建./images目录并放入至少3张不同类别的图片（如猫、狗、汽车各一张），将查询图片命名为query.jpg放在项目根目录（查询图片可以是图像库中的任意一张，用于验证检索正确性）。

运行：

python blip_retrieval.py

判定标准：

程序正常输出Top-K结果，无报错；
若查询图片来自图像库，则自身应出现在结果中且相似度最高（接近1.0）；
相似度分数均在0~1之间，且同类图片的相似度明显高于不同类图片。

关键说明

get_image_features返回的向量已做L2归一化，因此点积结果即为余弦相似度，取值范围[-1,1]；
模型首次运行会自动下载权重（约990MB），需保持网络畅通；
若图像库包含查询图自身，其相似度应为1.0（浮点误差范围内），这是验证特征提取正确性的关键指标。