向量数据库实战：选型、调优与落地~系列文章21：向量数据库在图像搜索引擎中的应用：以图搜图的完整实现
向量数据库在图像搜索引擎中的应用以图搜图的完整实现 本文是《向量数据库实战选型、调优与落地》专栏第 21 篇⏱️阅读时间约 13 分钟 开篇以图搜图的原理淘宝拍照搜同款、Google 图片搜索、Pinterest 视觉搜索——这些产品的背后都是图像向量搜索。┌─────────────────────────────────────────────────────────┐ │ 以图搜图原理 │ ├─────────────────────────────────────────────────────────┤ │ │ │ 用户上传一张照片 │ │ │ │ │ ▼ │ │ ┌──────────────┐ │ │ │ CNN/ViT 模型 │ ← ResNet、EfficientNet、ViT │ │ │ 提取图像特征 │ │ │ └──────┬───────┘ │ │ │ │ │ ▼ │ │ 图像向量 [0.23, -0.45, 0.67, ...] (2048维) │ │ │ │ │ ▼ │ │ ┌──────────────┐ │ │ │ 向量数据库 │ ← 在百万张商品图中搜索最相似的 │ │ │ 相似度检索 │ │ │ └──────┬───────┘ │ │ │ │ │ ▼ │ │ ️ 返回相似商品图片 Top-20 │ │ │ └─────────────────────────────────────────────────────────┘ 完整实现Step 1选择图像特征提取模型模型维度速度精度推荐场景ResNet-502048⚡⚡⚡⭐⭐⭐⭐通用图像搜索EfficientNet-B41792⚡⚡⭐⭐⭐⭐⭐高精度场景ViT-B/16768⚡⚡⭐⭐⭐⭐⭐最新架构CLIP ViT-L/14768⚡⚡⭐⭐⭐⭐跨模态搜索Step 2图像入库importtorchimporttorchvision.modelsasmodelsimporttorchvision.transformsastransformsfromPILimportImagefrompymilvusimportCollectionimportnumpyasnp# 加载 ResNet-50去掉最后的分类层modelmodels.resnet50(pretrainedTrue)modeltorch.nn.Sequential(*list(model.children())[:-1])# 去掉 FC 层model.eval()# 图像预处理transformtransforms.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]),])defextract_image_features(image_path):提取图像特征向量imageImage.open(image_path).convert(RGB)tensortransform(image).unsqueeze(0)withtorch.no_grad():featuresmodel(tensor)returnfeatures.squeeze().numpy()# 批量入库defindex_images(image_dir,collection,batch_size100):批量索引图像importos all_paths[]all_vectors[]forfilenameinos.listdir(image_dir):iffilename.endswith((.jpg,.png,.jpeg)):pathos.path.join(image_dir,filename)vecextract_image_features(path)all_paths.append(path)all_vectors.append(vec)iflen(all_vectors)batch_size:collection.insert([all_paths,np.array(all_vectors).tolist()])all_paths[]all_vectors[]# 插入剩余ifall_vectors:collection.insert([all_paths,np.array(all_vectors).tolist()])collection.flush()print(f索引完成共{collection.num_entities}张图片)Step 3以图搜图defsearch_similar_images(query_image_path,collection,top_k20):以图搜图# 1. 提取查询图片特征query_vecextract_image_features(query_image_path)# 2. 向量搜索collection.load()resultscollection.search(data[query_vec.tolist()],anns_fieldembedding,param{metric_type:L2,params:{ef:128}},limittop_k,output_fields[image_path])# 3. 返回结果similar_images[]forhitinresults[0]:similar_images.append({path:hit.entity.get(image_path),distance:hit.distance,similarity:1/(1hit.distance)# 转换为相似度})returnsimilar_images# 使用resultssearch_similar_images(query_photo.jpg,collection)forrinresults[:5]:print(f相似度:{r[similarity]:.4f}|{r[path]}) 图像搜索优化技巧优化项方法效果数据增强旋转、翻转、裁剪后取平均向量鲁棒性提升 15%特征归一化L2 归一化距离计算更稳定多尺度特征融合多层特征精度提升 5-10%GPU 加速特征提取用 GPU速度提升 10x缓存热门查询高频查询结果缓存延迟降低 50% 本篇核心要点回顾要点说明图像特征提取ResNet-50 / EfficientNet / ViT距离度量图像搜索推荐用 L2欧氏距离入库流程图片 → CNN → 特征向量 → 向量数据库搜索流程查询图片 → 特征提取 → 向量搜索 → 返回相似图性能关键特征提取是瓶颈用 GPU 加速下篇预告《生产环境避坑指南数据一致性、故障恢复、版本升级的血泪教训 》有问题欢迎评论区讨论觉得有用请点赞收藏 作者高炉炼铁智能化技术研究者专注钢铁冶金与人工智能 交叉领域。 如果觉得有帮助请点赞、收藏、转发版权归作者所有未经许可请勿抄袭套用商用(或其它具有利益性行为)。 关注专栏不错过后续精彩内容
)
