以图搜图：基于深度特征的商品与场景检索

bash pip install torch torchvision PIL numpy matplotlib faiss-cpu

- `torch` + `torchvision`：提供预训练模型和图片预处理。
- `Pillow`：图片读取与处理。
- `faiss-cpu`：向量相似度搜索（若需GPU加速可安装`faiss-gpu`）。
- `matplotlib`：结果可视化。

### 4. 特征提取：使用预训练ResNet

我们将使用在ImageNet上预训练的ResNet-50作为特征提取器，移除最后的分类层，取倒数第二层输出作为特征向量。这样提取的特征泛化能力强，适合通用图像检索。

```python
import torch
import torch.nn as nn
from torchvision import models, transforms
from PIL import Image

# 定义特征提取器
class FeatureExtractor(nn.Module):
    def __init__(self):
        super().__init__()
        base_model = models.resnet50(weights='ResNet50_Weights.DEFAULT')
        # 移除最后的全连接层和平均池化层，保留到layer4输出
        self.features = nn.Sequential(*list(base_model.children())[:-2])
        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))  # 全局平均池化

    def forward(self, x):
        x = self.features(x)
        x = self.avgpool(x)
        return torch.flatten(x, 1)  # 输出形状: (batch, 2048)

# 实例化模型并设为评估模式，不更新梯度
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = FeatureExtractor().to(device)
model.eval()

输入图片需要经过标准化预处理：

transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                         std=[0.229, 0.224, 0.225])
])

def extract_features(image_path):
    """单张图片特征提取"""
    img = Image.open(image_path).convert('RGB')
    img_tensor = transform(img).unsqueeze(0).to(device)  # 增加batch维度
    with torch.no_grad():
        features = model(img_tensor)
    return features.cpu().numpy().flatten()  # 转为numpy数组

这样，extract_features("bag.jpg") 会返回一个2048维的浮点向量。

5. 构建特征数据库并建立FAISS索引

假设我们的图库中有1000张商品图片。我们需要预计算所有图片的特征，并存入FAISS索引中，以便实时查询。

import faiss
import numpy as np
import os

# 准备图库图片路径列表
image_dir = 'product_images'
image_paths = [os.path.join(image_dir, f) for f in os.listdir(image_dir) if f.endswith(('.jpg', '.png'))]

# 批量提取特征
feature_list = []
for path in image_paths:
    feat = extract_features(path)
    feature_list.append(feat)

# 转为(n, 2048)的矩阵，注意类型为float32
feature_matrix = np.array(feature_list).astype('float32')

# 创建FAISS索引
dim = feature_matrix.shape[1]  # 特征维度
index = faiss.IndexFlatL2(dim)  # 使用L2欧氏距离
index.add(feature_matrix)       # 将特征矩阵添加到索引中

# 可选：保存索引和路径映射，以便后续使用
faiss.write_index(index, "image_index.faiss")
import pickle
with open("paths.pkl", "wb") as f:
    pickle.dump(image_paths, f)

此时，index 索引已经包含了全部图库特征，使用L2距离衡量相似度。对于商品检索，也可以尝试余弦相似度（FAISS中可先对向量做L2归一化，再用内积搜索，相当于余弦相似度）。

6. 执行相似图像查询

当用户输入一张查询图片时，我们提取其特征，然后从索引中找出最近的K张图。

def search(query_image_path, k=5):
    # 提取查询图特征
    query_feat = extract_features(query_image_path)
    query_feat = query_feat.astype('float32').reshape(1, -1)

    # 搜索最近邻
    distances, indices = index.search(query_feat, k)

    # 打印结果
    print("查询图片:", query_image_path)
    for i, idx in enumerate(indices[0]):
        print(f"Top {i+1}: {image_paths[idx]}  距离: {distances[0][i]:.4f}")
    return indices[0], distances[0]

# 测试
result_indices, _ = search("query_bag.jpg", k=5)

7. 可视化搜索结果

用matplotlib并排展示查询图和Top-K结果，直观感受效果。

import matplotlib.pyplot as plt

def show_results(query_path, result_indices, num_results=5):
    fig, axes = plt.subplots(1, 1+num_results, figsize=(15, 5))
    axes[0].imshow(Image.open(query_path))
    axes[0].set_title("Query")
    axes[0].axis('off')

    for i, idx in enumerate(result_indices[:num_results]):
        img = Image.open(image_paths[idx])
        axes[i+1].imshow(img)
        axes[i+1].set_title(f"#{i+1}")
        axes[i+1].axis('off')
    plt.tight_layout()
    plt.show()

show_results("query_bag.jpg", result_indices)