☰
细粒度图像检索实战:从特征抽取到FAISS索引的完整方案
2026/10/5 7:28:59 网站建设 项目流程

简介:这是基于Python构建的细粒度图像检索系统完整源码包,面向计算机视觉研究者、算法工程师及需要技术汇报的开发者,可用于解决同类商品、生物特征等细粒度对象的精确检索问题。包内共79个文件、约66.21MB,其中37个py脚本实现三元组损失网络、多标签学习、SIFT词包检索等核心算法,15个txt与11个pdf提供数据说明与研究论文,4个pptx和docx辅助项目汇报,3个pyc、keep等文件完善工程结构。目前已有349人学习下载。这套资源不仅包含可运行的训练与测试代码,还整理了细粒度属性学习、深度哈希检索等方法的阅读材料,并配套CUB、VOC2012等数据集索引,便于研究者快速复现实验、梳理汇报素材,是系统理解图像检索技术链条的实用参考资料。

1. 细粒度图像检索:为什么常规以图搜图在这里会失效

细粒度图像检索(Fine-Grained Image Retrieval,FGIR)要解决的问题不是“找同类物体”,而是“在同一大类里找出同一个子类”。拿最常见的 CUB-200-2011 鸟类数据集举例:两百种鸟在普通人眼里都是“麻雀大小、羽毛灰褐色”,但系统要能在这种几乎不可见的差异里,把查询图的鸟和数据库里同一物种的图排到最前面。常规的以图搜图系统,用 ImageNet 预训练模型抽取全局特征就能工作得很好,但到了细粒度场景,数量级接近的分类间差异被全局池化抹平,top10 全部翻车是常态。这类系统的落地场景集中在生态监测、商品局部瑕疵溯源、医疗影像子型分类、植物标本归档等方向,适合手里有一批“类内差异小、类间差异更小”的图像数据、又不想人工逐张标注的从业者。文本后面讲的方案不需要分布式集群,一台带 CUDA 的 GPU 机器加 Python 就能完整跑通。

2. 数据准备:细粒度检索的第一步是明确“查询集”和“底库集”

2.1 为什么拿 CUB-200-2011 作为首个验证集

做细粒度图像检索,最先要面对的是数据集选择。常见的选择有 CUB-200-2011(鸟类)、Stanford Cars(车型)、Oxford Flowers 102(花卉),三者都是学术界的公开基准,其中 CUB 因为类别数多、类间差异小,是检验 FGIR 方案是否有效的“试金石”。CUB-200-2011 包含约 6000 张训练图和 5800 张测试图,分属 200 个类别,每类约 30 张训练图,这个数据量决定了你不能像训练 ImageNet 那样从零训练深度网络,必须走迁移学习的路子。

另外一个关键点是“查询集”和“底库集”必须从数据集层面就拆开。很多初次接触检索系统的人把训练集直接当作底库,拿测试集去查,结果评估指标虚高,上线后一查一个错。通常的做法是:底库使用测试集(或者单独划分的 gallery 集),查询集再另外划分一 份 query 集,两份数据严格不相交。这样评估出来的指标才代表真实线上表现。

2.2 自定义 Dataset:把目录结构、标签和图像路径串起来

PyTorch 的torchvision.datasets.ImageFolder能处理按类别分目录的数据,但细粒度数据集经常带额外元数据,比如 CUB 的 image_id、类别属性、bounding box 标注。最省事的做法是直接用ImageFolder读图,配合一个id2label映射文件一起用。我一般会把数据集按下面的目录结构整理:

cub200/ images/ 001.Black_footed_Albatross/ Black_Footed_Albatross_0001_796111.jpg ... 002.Laysan_Albatross/ ... train_test_split.txt classes.txt

整理好之后,写一个继承torch.utils.data.Dataset的加载器。下面是实测可用的最小实现:

import os from PIL import Image from torch.utils.data import Dataset class FineGrainedDataset(Dataset): def __init__(self, root, split_file, transform=None): """ root: 图像根目录,如 cub200/images split_file: CUB 官方 train_test_split.txt, 格式为 "image_id is_train_image" """ self.root = root self.transform = transform self.image_paths = [] self.labels = [] # 读取划分文件,1 表示训练图,0 表示测试图 with open(split_file, 'r') as f: lines = f.readlines()[1:] # 第一行是表头,跳过 for line in lines: parts = line.strip().split() if len(parts) < 2: continue img_id = int(parts[0]) is_train = int(parts[1]) if is_train == 1: # 这里控制加载训练集 # CUB 的 image_id 从 1 开始,和文件名一一对应 folder = os.path.join(root, self._find_folder(img_id)) self.image_paths.append( os.path.join(root, folder, f"{img_id:04d}.jpg") ) # 类别标签从目录名前缀解析,如 001.Black_footed_Albatross label = int(folder.split('.')[0]) - 1 self.labels.append(label) def _find_folder(self, img_id): # 实际项目中可以用一个预生成的 id2folder 字典,避免每次都扫目录 for folder in os.listdir(self.root): for fname in os.listdir(os.path.join(self.root, folder)): if fname.startswith(f"{img_id:04d}_"): return folder raise FileNotFoundError(f"image {img_id} not found") def __len__(self): return len(self.image_paths) def __getitem__(self, idx): path = self.image_paths[idx] image = Image.open(path).convert('RGB') label = self.labels[idx] if self.transform: image = self.transform(image) return image, label

这段代码的逻辑不复杂:从 CUB 的train_test_split.txt读取每张图是否属于训练集,然后根据 image_id 去 images 目录下找到对应文件。_find_folder方法先用最粗暴的方式遍历目录——100 多行代码能跑通,但每次构建 Dataset 都会全盘扫描一次,200 类大约几万个文件,初始化会慢几秒。工程上更常见的做法是提前把image_id -> 文件夹名的关系用字典保存成 JSON,启动时一次性加载。

2.3 预处理流水线:分辨率、归一化和数据增强的选择

细粒度检索的预处理有三个参数直接决定效果。第一个是输入分辨率。ResNet 系列在 ImageNet 上默认用 224×224,但细粒度特征往往藏在羽毛纹理、车灯轮廓这些局部区域,224 会丢失太多细节。常见做法是训练时用随机裁剪 224 或 256 作为数据增强,提取检索特征时统一用 448×448 或至少 384×384 的原图中心裁剪。增大分辨率带来的收益非常直接,但显存和推理耗时也会同步上涨,需要做取舍。

第二个是归一化参数。ImageNet 预训练模型要求输入按mean=[0.485, 0.456, 0.406]、std=[0.229, 0.224, 0.225]做标准化,迁移到细粒度数据集时这个参数不要改。有些同学在微调时为了使输入“适应自己的数据”重新统计均值和方差,结果特征分布偏移,检索精度反而掉三四个点。第三个是数据增强,训练阶段用 RandomResizedCrop 加 RandomHorizontalFlip 就够,不要上 AutoAugment 这种重增强——细粒度类别差异本来就细微,过度增强会把关键判别区域随机裁掉,模型学到的是背景统计量而不是鸟类特征。

3. 特征抽取:迁移学习与 Embedding 层的设计

3.1 选 ResNet50 还是 Vision Transformer

细粒度检索的特征抽取器,核心原则只有一个:用在大规模数据集上预训练过的模型,不要在百万级以下的数据集上从头训练。当前从业者最常用的两个选择是 ResNet50 和 ViT-B/16。ResNet50 的优点是显存占用小、推理快、在中小数据集上不容易过拟合,微调时只需替换最后的全连接层;ViT 的特征在细粒度任务上往往更准,但需要更大的训练数据支撑,在 CUB 这种每类只有三十张图的场景下很快就过拟合。

我一般会先在 ImageNet 预训练 ResNet50 上微调一个分类头做 warm-up,然后再把分类头替换成固定长度的 embedding 层。ResNet50 的avgpool输出是 2048 维,这个维度作为检索特征已经够用。你也可以在 avgpool 之后接一个nn.Linear(2048, 512)把特征压到 512 维,好处是底库规模大时索引占用内存更小、检索速度更快,代价是精度会有一点点损失。

3.2 抽取特征并落盘的实现脚本

下面这段脚本做的事情是加载训练好的模型,去掉最后的分类层,对数据集里的每一张图像做前向推理,把归一化后的特征向量保存成.npy文件。

import torch import torchvision.transforms as T import numpy as np from torchvision import models from PIL import Image from tqdm import tqdm device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') # 使用 ImageNet 预训练权重,替换分类头为恒等映射 model = models.resnet50(weights=models.ResNet50_Weights.IMAGENET1K_V1) model.fc = torch.nn.Identity() # 直接输出 2048 维特征 model = model.to(device) model.eval() # 448x448 中心裁剪,比训练时的 224 保留更多细节 transform = T.Compose([ T.Resize(512), T.CenterCrop(448), T.ToTensor(), T.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 = transform(img).unsqueeze(0).to(device) with torch.no_grad(): feat = model(img) # shape: [1, 2048] feat = feat.squeeze(0).cpu().numpy() # L2 归一化:让特征只保留方向信息,去除图像亮度影响 feat = feat / np.linalg.norm(feat) return feat # 假设 image_paths 是前面 Dataset 里收集的全部路径 features = [] for p in tqdm(image_paths): feat = extract_features(p) features.append(feat) features = np.stack(features).astype('float32') # [N, 2048] np.save('gallery_features.npy', features) with open('gallery_paths.txt', 'w') as f: for p in image_paths: f.write(p + '\n')

model.fc = torch.nn.Identity()是这段代码的关键操作,它把全连接分类层替换成恒等映射,模型输出的就是 avgpool 后的 2048 维特征向量。L2 归一化这一步容易被省略,但它决定了后续是用余弦相似度还是欧式距离。归一化之后,两个特征的余弦相似度等价于归一化向量的内积,同时避免了不同图像整体亮度、对比度差异对特征尺度的影响。

3.3 为什么不直接拿 softmax 概率当检索特征

不少初次接触细粒度检索的工程师,会把分类模型最后一层 softmax 输出的 200 维概率向量直接当作检索特征,理由是“分到同一类的图概率分布应该接近”。这个做法在类别数少时还能工作,类别一多就崩。原因在于 softmax 概率受训练类别分布影响,分类头学到的决策边界是按照“区分所有训练类”的目标优化的,而检索要求的是“相似物体在特征空间距离近”,这两个目标不完全一致。用分类概率做特征,结果就是底库里同一物种的图像可能因为分类头把某张图误判到相近类别,导致概率分布偏移,检索排序跟着出错。

正确做法是把分类头只当成训练阶段的辅助监督信号,推理阶段丢弃它,使用 penultimate layer 的特征。另外,如果想让特征更适配检索任务,可以用三元组损失(Triplet Loss)或 ArcFace 这类度量学习损失在分类预训练的基础上继续微调 embedding 层。三元组选择的策略优先挑“同类最难样本”和“异类最近样本”,这能显著拉近同子类图像在特征空间的距离。

4. 索引与检索:从暴力遍历到底库千万级的方案选型

4.1 相似度度量与特征归一化

特征抽取完成之后,系统就变成标准的向量检索问题了。细粒度图像检索里最常见的相似度度量是余弦相似度,计算方式是两个 2048 维向量归一化后的内积。为什么不用欧式距离?因为特征向量不同维度的尺度差异很大,欧式距离会放大数值较大的维度对结果的影响,而归一化内积天然地压缩了这个问题。

底库规模决定索引方案。几百到一万张图,直接用 numpy 暴力检索就够了;十万级以上,才需要引入 FAISS 或 Milvus。暴力检索最大的优势是精确召回——它能保证返回全局最近的 K 个向量,不做任何近似,这对细粒度检索这种“差之毫厘就判错类”的场景非常重要。

4.2 numpy 暴力检索:十行代码跑通精确召回

下面的代码演示一个查询图的特征如何与底库所有特征做相似度计算并返回 topK 结果。

import numpy as np # 假设 gallery_features 是 (N, 2048) 的底库特征数组 gallery_features = np.load('gallery_features.npy') # 已经 L2 归一化 gallery_paths = open('gallery_paths.txt').read().splitlines() def search_similar(query_feat, top_k=10): """ query_feat: 查询图的 2048 维 L2 归一化特征 返回 (topk_index, topk_score) """ # 查询特征也要归一化 query_feat = query_feat.reshape(1, -1).astype('float32') query_feat = query_feat / np.linalg.norm(query_feat, axis=1, keepdims=True) # 归一化向量内积 == 余弦相似度 scores = np.dot(gallery_features, query_feat.T).squeeze() # 按相似度降序,返回前 top_k 个索引 topk_idx = np.argsort(scores)[::-1][:top_k] topk_score = scores[topk_idx] return topk_idx, topk_score query_idx, query_score = search_similar(query_feat, top_k=10) for i, score in zip(query_idx, query_score): print(f"{gallery_paths[i]:60s} score={score:.4f}")

这段代码有几个容易踩的细节。一是np.argsort默认升序,必须加[::-1]反转。二是内存占用,gallery_features是 N 行 2048 列的 float32 数组,一万张图占用约 80MB,在内存里做矩阵乘法很快;但如果 N 到百万级别,np.dot一次要算 100 万×2048 维的浮点运算,响应时间会到几百毫秒甚至秒级。三是查询特征的归一化不要漏,否则结果会偏好亮度高的图像。

4.3 FAISS 索引:万级底库的性能分水岭

当底库图像数量超过五万张,暴力检索的延迟就变得不可接受了。Facebook 的 FAISS 是当前业内最通用的向量检索库,支持精确检索和多种近似检索索引。精确索引IndexFlatIP的召回率和 numpy 暴力检索完全一致,只是底层用 BLAS 优化过,速度更快;近似索引IndexIVFFlat则把底库聚类成 nlist 个桶,查询时只搜索最近的 nprobe 个桶,性能提升明显但召回率有所下降。

先安装依赖:

pip install faiss-cpu numpy

FAISS 构建索引和查询的示例代码如下:

import faiss import numpy as np # gallery_features 已经 L2 归一化,shape (N, 2048) gallery_features = np.load('gallery_features.npy').astype('float32') dim = gallery_features.shape[1] nlist = 100 # 聚成 100 个桶 nprobe = 8 # 查询时搜索的桶数量 index = faiss.IndexIVFFlat(faiss.IndexFlatIP(dim), dim, nlist) # IndexIVFFlat 必须先训练(对底库聚类),再添加向量 index.train(gallery_features) index.add(gallery_features) index.nprobe = nprobe query_feat = query_feat.astype('float32').reshape(1, -1) # 返回 topK 的相似度分数和索引,score 越大越相似 scores, pred_idx = index.search(query_feat, top_k) for s, idx in zip(scores[0], pred_idx[0]): print(gallery_paths[idx], s)

参数nlist与nprobe直接影响速度与召回率的平衡。nlist越大,每个桶里向量越少,但训练聚类耗时越长;nprobe越大,检索越接近精确结果,但耗时近似线性上升。做工程落地时我一般先固定nlist=100,用nprobe从 1 到 16 扫一遍,选一个 mAP 下降不超过 1 个百分点的最小值。如果你追求全量精确召回,IndexFlatIP永远是最稳妥的选择;例如底库 10 万张图,IndexFlatIP单次查询延迟在几十毫秒量级,完全够用。

5. 细粒度检索系统开发中的 5 个坑:现象、原因和修复办法

5.1 训练集与底库集数据泄漏,评估指标虚高

现象:离线测试时 mAP 高达 0.85,部署到线上后效果断崖下跌,用户反馈检索结果不相关。原因:把训练集图像塞进了底库,查询集虽然独立,但训练过程中模型已经见过底库图像,特征空间被强行拉近了。解决:从数据集划分开始就固定三份独立数据——训练集用于更新模型参数,底库集只用于构建索引,查询集只用于评估。三者不能有交集,哪怕同一张图的不同裁剪版本也算泄漏,因为模型已经见过像素信息。

5.2 特征没有归一化,检索结果偏向高亮度图像

现象:返回结果里白色羽毛的鸟明显多于深色羽毛的鸟,即使类别不同。原因:特征向量的数值范围受图像整体亮度影响,未归一化的特征在计算内积时,亮度高的图像特征向量模长更大,天然获得更高分数。解决:在抽取特征后、构建索引前,对每个特征向量做 L2 归一化,并且查询特征做同样的归一化。这条要写成数据管线的固定步骤,而不是靠每个人手动记住。

5.3 用分类准确率评估检索模型

现象:分类准确率从 78% 涨到 82%,但检索的 top10 命中率几乎没有变化。原因:分类和检索是两种任务,分类只关心能否预测对,检索关心的是相似样本能否聚集在相近的位置。分类准确率高说明模型学到了类别的判别信息,但这些信息可能集中在少数几个神经元上,特征的支撑集不稳定,检索时稍有扰动排序就变了。解决:单独用 mAP、R@K 等检索指标评估模型,这两类指标才是检索系统的 KPI。分类准确率只能作为训练曲线的参考。

5.4 局部特征不足,仅靠全局特征压不住混淆类

现象:两个亚种外形极其相似,全局特征距离很近,top5 里混入大量近亲类别。原因:avgpool对整张图做平均池化,把细小的判别性局部区域(比如鸟嘴的形状差异)平均掉了。解决:有两条路可走,简单的一条是提高输入分辨率到 448 或 512,让局部纹理进入特征;另一条是用多尺度特征融合,把 ResNet 最后一个 block 的 feature map 按空间网格做区域池化,生成多个局部特征,检索时取全局特征和局部特征得分的加权和。

5.5 服务器端推理环境与训练环境不一致

现象:训练好的模型在本地测试正常,部署到 GPU 服务器后特征变了,检索结果不可复现。原因:常见的有三种,一是ResNet50默认有 BatchNorm 层,模型在train()和eval()模式下行为不一致,部署时忘记调用model.eval();二是推理时用的图像预处理与训练时不一致,比如缩放尺寸、中心裁剪位置有偏差;三是混合精度推理时特征没有转换回 float32。解决:写一个推理自检脚本,固定几张图,对比训练环境和部署环境的特征向量余弦相似度,低于 0.999 就逐项排查上述三个原因。

6. 检索效果验证的正确姿势与重排序技巧

6.1 R@K 与 mAP:比肉眼看图更可靠的评估方法

做细粒度检索系统,最忌拿着几张查询图“目测”返回结果不错就宣布上线。肉眼只能看个大概,遇到边界混淆案例就分不清是算法问题还是数据问题。我每次改完特征抽取或索引参数,都会先跑一遍离线评估脚本,算 R@K 和 mAP,用数字说话。R@K 表示查询图的前 K 个返回结果里有多少比例包含与查询图同类的图像;mAP 则是所有查询的 Average Precision 均值,对排序位置敏感,更适合衡量整体排序质量。

下面这个评估脚本只需要底库特征、底库标签、查询特征、查询标签四样输入,就能在几秒内给出全部指标。

import numpy as np def evaluate_retrieval(query_feats, query_labels, gallery_feats, gallery_labels, K=10): """ query_feats / gallery_feats: L2 归一化特征 query_labels / gallery_labels: 类别 id,来自数据集标注 """ # 所有查询与底库的相似度矩阵 sim_matrix = np.dot(query_feats, gallery_feats.T) # [Q, N] ap_sum = 0.0 rks = [] for q_idx in range(query_feats.shape[0]): q_label = query_labels[q_idx] sorted_idx = np.argsort(sim_matrix[q_idx])[::-1] # 降序 # R@K: 前 K 个里是否存在同类 rk_true = any(gallery_labels[idx] == q_label for idx in sorted_idx[:K]) rks.append(rk_true) # AP: 用累计位置计算精确率均值 hits = 0 prec_sum = 0.0 for rank, idx in enumerate(sorted_idx): if gallery_labels[idx] == q_label: hits += 1 prec_sum += hits / (rank + 1) ap = prec_sum / max(np.sum(gallery_labels == q_label), 1) ap_sum += ap recall_at_k = np.mean(rks) mAP = ap_sum / query_feats.shape[0] return recall_at_k, mAP r_at_k, map_score = evaluate_retrieval( query_feats, query_labels, gallery_feats, gallery_labels, K=10 ) print(f"R@10 = {r_at_k:.4f}, mAP = {map_score:.4f}")

这个脚本刻意用 numpy 实现,不依赖任何重框架。底库特征规模一万条以内,跑一遍只需要几秒。实际做实验结果对比时,我会把不同分辨率、不同特征维度、不同索引参数的结果列在同一张表里,R@10 看“检索到底有没有找到对的”,mAP 看“对的图排得靠不靠前”。两者要同时看,R@10 高 mAP 低,说明系统能找到同类但排序混乱,优化重点在重排。

6.2 重排序:用空间匹配给相似度“验真”

即便特征抽取做得很到位,纯全局特征的 topK 结果里仍然会混入少数外形相似的异常样本。一个工程上性价比较高的技巧是二级重排序:第一轮用 FAISS 或暴力检索拿到 top50 候选,第二轮把查询图和每个候选图分别切成 2×2 或 3×3 的局部块,分别抽取局部特征,计算查询图和候选图每个对应位置子块之间的特征相似度,取各区域相似度的加权平均值作为新的排序分数。

# 伪代码描述重排流程,实际实现时需要复用前面的特征抽取函数 candidates = faiss_search(query_feat, top_k=50) regions = [(0,0), (0,1), (1,0), (1,1)] # 横纵等分 for cand in candidates: cand_regions = extract_grid_features(cand.path, grid=(2,2)) q_regions = extract_grid_features(query_path, grid=(2,2)) score = mean([cos_sim(q, c) for q, c in zip(q_regions, cand_regions)]) cand.rerank_score = score

这里有两点要注意。一是局部块的大小不要切得太碎,2×2 或 3×3 已经足够捕获得细粒度判别区域,再细会把同一物体的不同部件错误强绑定。二是重排分数和第一轮全局分数做加权融合时,权重通常取全局 0.7、局部 0.3,以全局为主、局部拉偏。哪边权重过高都会让结果在一部分查询上退化。

6.3 检索结果可视化:横坐标标签太密集的处理

调试检索系统时,把查询图和 top10 结果按序排列画成一行对比图是最直观的手段。我第一次画这种图时用的是 Matplotlib 直接用类别名当 x 轴标签,200 个类别挤在一起,横坐标密密麻麻完全分不清谁是谁,调试效率极低。后来改用两行式画法:第一行按编号 1~10 标注候选图顺序,第二行在每张图下方单独写它的类别名,只显示前 10 个候选的类别,其余省略号缩写。这样既保留了完整类别信息,又避免了横坐标太密集的问题。可视化工具不挑贵贱,能把 badcase 快速挑出来就行。

做细粒度检索项目一年下来,我最大的体感教训是:不要迷信某个新提出的特征抽取模块能带来多少个点的提升,先把数据划分、特征归一化、评估脚本这些地基打扎实。地基歪了,后面算法怎么换都是白搭。另一个收尾习惯是每个实验跑完都存档完整的特征文件和参数表,否则两三天后连自己都分不清哪份特征用的是 448 分辨率还是 512 分辨率。这条经验是从小项目到大项目一路踩过来的,希望帮到你。

本文还有配套的精品资源,点击获取

需要专业的网站建设服务?

联系我们获取免费的网站建设咨询和方案报价,让我们帮助您实现业务目标

立即咨询