简介这份资源面向人工智能与信息检索方向的学习者提供一套基于VGG-16的图像检索系统完整项目实践。项目以深度学习特征替代传统颜色、形状、纹理等手工特征采用Keras预训练模型对图像库逐张抽取特征并存入h5文件检索阶段以夹角余弦度量相似度对测试图片特征与库中向量逐一计算距离并排序返回相似度Top3结果覆盖特征提取、索引化数据集、相似度定义与检索全流程。资源包共253个文件以241张jpg图像样本为主另含5个xml标注、3个py脚本、2个pyc编译文件、1个h5特征库及1个iml工程配置压缩包约41.13MB目录结构便于按模块查阅。目前已有197人学习下载适合希望理解图像检索原理、动手复现特征提取与相似度排序流程的读者参考借鉴。1. 基于VGG-16的图像检索系统从特征提取到相似度排序的完整落地路径图像检索系统要解决的核心问题很具体用户给一张查询图系统从图库里找出视觉上最相似的若干张。传统做法靠人工标注关键词成本高且覆盖不全而基于VGG-16的方案本质是把图像映射成一个高维特征向量再用向量距离衡量相似度。VGG-16之所以常被选作骨干是因为它的卷积层堆叠结构规整、预训练权重容易获取、提取出的全连接层特征通常是fc2的4096维在通用图像上泛化能力稳定。这套方案适合谁适合手头有几千到几十万张图、想做以图搜图原型、又不想从头训练模型的开发者。下面从原理到代码把这条链路拆开讲清楚。2. VGG-16为什么能当图像检索的特征提取器2.1 从分类网络到特征编码器的迁移逻辑VGG-16最初是为ImageNet分类任务设计的输入224×224的RGB图像经过13个卷积层和3个全连接层最后输出1000类概率。做图像检索时我们不需要它的分类结果而是取倒数第二层全连接层fc2的输出——一个4096维向量。这个向量可以理解为图像在VGG-16学到的语义空间中的坐标。为什么这个坐标能用来检索因为卷积层提取的是局部纹理和边缘越往后的全连接层越偏向整体语义。fc2的4096维已经足够表达“这是一只狗”或“这是一辆车”这种高层概念同时保留了类内差异。相比直接用像素做欧氏距离VGG-16特征对光照、轻微旋转、尺度变化更鲁棒。常见做法是去掉最后的分类层把网络当作一个固定的特征提取器不参与训练。这意味着你不需要标注数据也不需要GPU集群一张一张图前向传播即可。对于中小规模图库这是性价比最高的起点。2.2 特征提取的完整代码实现下面这段代码用PyTorch加载预训练VGG-16去掉分类头对单张图像提取4096维特征。注意输入预处理必须和ImageNet训练时一致。import torch import torchvision.models as models import torchvision.transforms as transforms from PIL import Image # 加载预训练VGG-16weights参数指定使用ImageNet权重 model models.vgg16(weightsmodels.VGG16_Weights.IMAGENET1K_V1) # 只保留features和avgpool以及classifier的前两层到fc2为止 # classifier结构: Linear(25088,4096) - ReLU - Dropout - Linear(4096,4096) - ReLU - Dropout - Linear(4096,1000) feature_extractor torch.nn.Sequential( model.features, model.avgpool, torch.nn.Flatten(), model.classifier[0], # fc1: 25088 - 4096 model.classifier[1], # ReLU model.classifier[3], # fc2: 4096 - 4096 model.classifier[4], # ReLU ) # 切换到推理模式关闭Dropout feature_extractor.eval() # 预处理Resize到256CenterCrop到224转Tensor按ImageNet均值方差归一化 preprocess transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]), ]) def extract_feature(img_path): img Image.open(img_path).convert(RGB) tensor preprocess(img).unsqueeze(0) # 增加batch维度 with torch.no_grad(): feat feature_extractor(tensor) return feat.squeeze(0).numpy() # 返回4096维numpy数组逻辑说明model.features输出7×7×512的特征图avgpool变成7×7Flatten拉成25088维再经过fc1和fc2得到4096维。参数上weights指定IMAGENET1K_V1表示使用官方预训练权重Resize(256)和CenterCrop(224)是VGG-16的标准输入流程不要随意改成直接Resize到224否则会损失边缘信息。torch.no_grad()关闭梯度计算节省显存并加速。2.3 批量提取与特征库构建单张提取太慢实际项目要批量处理整个图库。下面用DataLoader做批处理并把结果存成numpy文件。import os import numpy as np from torch.utils.data import DataLoader, Dataset class ImageFolderDataset(Dataset): def __init__(self, root_dir, transform): self.root_dir root_dir self.transform transform self.img_names sorted(os.listdir(root_dir)) self.img_paths [os.path.join(root_dir, n) for n in self.img_names] def __len__(self): return len(self.img_paths) def __getitem__(self, idx): img Image.open(self.img_paths[idx]).convert(RGB) return self.transform(img), idx dataset ImageFolderDataset(gallery/, preprocess) loader DataLoader(dataset, batch_size32, shuffleFalse, num_workers4) all_feats [] all_indices [] with torch.no_grad(): for batch_imgs, batch_idx in loader: feats feature_extractor(batch_imgs) all_feats.append(feats.numpy()) all_indices.extend(batch_idx.numpy()) feature_matrix np.vstack(all_feats) # shape: (N, 4096) np.save(gallery_features.npy, feature_matrix) np.save(gallery_indices.npy, np.array(all_indices))参数说明batch_size32在8GB显存下比较稳妥如果显存不够降到16或8num_workers4根据CPU核数调整Windows下建议设为0避免多进程问题。feature_matrix的每一行对应一张图顺序由gallery_indices记录后续检索时靠这个索引找回原图路径。3. 相似度计算与检索排序的工程实现3.1 余弦相似度 vs 欧氏距离选哪个特征向量比较常用两种度量余弦相似度和欧氏距离。VGG-16的fc2输出经过ReLU所有值非负向量方向比绝对长度更有意义。余弦相似度只看向量夹角对整体亮度变化不敏感通常效果更稳。欧氏距离受向量模长影响大如果图库中某些图特征值普遍偏大会主导排序结果。我一般先把特征向量做L2归一化然后直接用内积计算余弦相似度。归一化后内积等价于余弦值且可以用矩阵乘法一次性算出查询图与所有图库图的相似度速度极快。# 对特征库做L2归一化 norms np.linalg.norm(feature_matrix, axis1, keepdimsTrue) feature_matrix_norm feature_matrix / (norms 1e-10) def search(query_feat, top_k10): # query_feat: 4096维numpy数组 q_norm query_feat / (np.linalg.norm(query_feat) 1e-10) # 内积得到余弦相似度 sims feature_matrix_norm.dot(q_norm) # shape: (N,) # 取top_k索引 top_indices np.argsort(sims)[::-1][:top_k] return top_indices, sims[top_indices]逻辑说明np.argsort(sims)[::-1]先升序排序再反转得到降序索引。top_k一般设10到20太多会引入不相关结果。1e-10防止除零。如果图库超过百万级dot仍然很快但内存占用要考虑可以改用faiss做近似最近邻。3.2 用faiss加速大规模检索当图库到十万级以上暴力内积会变慢。faiss是Facebook开源的向量检索库支持GPU和多种索引结构。下面用IndexFlatIP内积索引做精确检索适合百万级以下。import faiss # 构建内积索引维度4096 index faiss.IndexFlatIP(4096) # 添加归一化后的特征 index.add(feature_matrix_norm.astype(float32)) def search_faiss(query_feat, top_k10): q_norm query_feat / (np.linalg.norm(query_feat) 1e-10) q q_norm.astype(float32).reshape(1, -1) sims, indices index.search(q, top_k) return indices[0], sims[0]参数说明IndexFlatIP是精确内积索引召回率100%但内存占用为N×4096×4字节。如果图库到千万级改用IndexIVFFlat并设置nlist聚类中心数一般取sqrt(N)检索时指定nprobe探测簇数越大越准但越慢。faiss要求输入float32注意类型转换。3.3 检索结果的可视化与评估检索完要验证效果。最直接的方法是拼图左边查询图右边top_k结果。同时可以算mAPmean Average Precision来量化前提是图库有类别标签。import matplotlib.pyplot as plt def show_results(query_path, result_indices, img_names): fig, axes plt.subplots(1, len(result_indices)1, figsize(3*(len(result_indices)1), 3)) axes[0].imshow(Image.open(query_path)) axes[0].set_title(Query) axes[0].axis(off) for ax, idx in zip(axes[1:], result_indices): img Image.open(os.path.join(gallery/, img_names[idx])) ax.imshow(img) ax.set_title(fRank {list(result_indices).index(idx)1}) ax.axis(off) plt.show()评估时如果查询图和图库中同类别图算相关可以计算每个查询的APAverage Precision再对所有查询取平均得到mAP。这一步能帮你判断特征提取和相似度度量是否合理而不是凭肉眼感觉。4. 避坑与常见问题排查4.1 现象检索结果全是同一类图但查询图明显不同原因特征向量没有归一化或者图库中某些类别的特征模长普遍偏大导致内积被模长主导。解决对图库和查询特征都做L2归一化确保相似度只反映方向差异。检查代码中是否遗漏了feature_matrix / norms这一步。4.2 现象提取特征时显存溢出CUDA out of memory原因batch_size太大或者没有用torch.no_grad()导致中间激活值保留计算图。解决把batch_size降到8或4确认with torch.no_grad():包裹了前向传播。如果仍然不够把模型放到CPU上跑速度慢但显存无压力。4.3 现象检索速度随图库增长线性下降十万张图要等好几秒原因用了暴力内积且没有用faiss或矩阵运算优化。解决改用faiss的IndexFlatIP或者至少把特征矩阵转成float32并用numpy.dot一次性算完。如果图库超过百万上IndexIVFFlatnlist设为int(sqrt(N))nprobe从10开始调。4.4 现象查询图和图库中同一张图完全一样的相似度不是1.0原因预处理不一致。查询图可能被Resize了两次或者归一化参数用错。解决确保查询图和图库图走完全相同的preprocess流程包括Resize、CenterCrop、ToTensor、Normalize。检查mean和std是否写成了[0.5,0.5,0.5]这种错误值。4.5 现象换用不同批次的图库检索效果波动大原因VGG-16的fc2特征对图像整体亮度、对比度敏感如果图库来源多样比如白天和夜间混在一起特征分布会漂移。解决在预处理中加入直方图均衡化或者改用对光照更鲁棒的骨干如ResNet。另一个做法是对特征做PCA白化降低各维度相关性。5. 进阶技巧用PCA降维和重排序提升检索质量5.1 PCA把4096维压到256维速度翻倍且精度损失很小4096维特征存储和计算成本高但很多维度是冗余的。用PCA降到256维检索速度提升十几倍mAP通常只掉1到2个点。做法是在图库特征上拟合PCA然后变换查询特征。from sklearn.decomposition import PCA pca PCA(n_components256, whitenTrue) feature_pca pca.fit_transform(feature_matrix_norm) # 图库降维 # 查询时 query_pca pca.transform(q_norm.reshape(1, -1))whitenTrue让各维度方差为1进一步消除尺度差异。注意PCA必须在图库特征上拟合不能单独对查询拟合。降维后重新建faiss索引维度改为256。5.2 重排序用查询扩展Query Expansion拉回漏掉的结果初次检索的top_k可能漏掉一些视觉上相似但特征距离稍远的图。查询扩展的思路是取初次检索的top_m比如top_5特征的平均值作为新的查询向量再检索一次。这样能把与多个近邻都相似的图排上来。def query_expansion(query_feat, top_m5, top_k10): indices, _ search_faiss(query_feat, top_m) # 取top_m特征的平均 expanded feature_matrix_norm[indices].mean(axis0) expanded expanded / (np.linalg.norm(expanded) 1e-10) return search_faiss(expanded, top_k)这个技巧在Oxford和Paris数据集上通常能提升3到5个点的mAP。代价是检索时间翻倍因为要算两次。如果延迟敏感可以只对top_20做重排序。5.3 一个我踩过的坑不要用fc1特征fc1是25088维比fc2的4096维大得多直觉上信息更丰富。但实测中fc1包含更多低级纹理对语义相似度反而不如fc2。我曾在花卉数据集上对比fc2的mAP比fc1高8个点。所以除非有特殊需求直接用fc2。5.4 验证方法用mAP和Top-1准确率双指标不要只看Top-1。Top-1只能反映最相似的那张对不对mAP能反映整个排序列表的质量。建议每次调整参数PCA维度、top_k、是否查询扩展后都在固定测试集上算这两个指标。如果mAP涨了但Top-1没动说明排序整体变好但头部没变如果Top-1涨了但mAP跌了可能是过拟合到某些简单样本。我自己的习惯是先跑通暴力检索拿到baseline mAP再逐步加PCA、faiss、查询扩展每加一个模块记录指标变化。如果某个模块让mAP掉了超过2个点就回退。这套流程帮我在多个图像检索项目里稳定交付希望帮到你。本文还有配套的精品资源点击获取
