从零搭建一个AI应用：用Python+Milvus快速构建你的第一个图像检索系统-编程实验室

从零搭建AI图像检索系统：Python与Milvus的实战指南

当你面对数千张未经分类的图片时，如何快速找到与某张图片内容相似的其他图片？传统的关键词搜索在这里完全失效，而基于深度学习的图像检索技术可以完美解决这个问题。本文将带你用Python和Milvus向量数据库，从零开始构建一个高效的图像相似度检索系统。

1. 系统架构与技术选型

一个完整的图像检索系统通常包含三个核心组件：

特征提取模型：将图片转换为高维向量表示
向量数据库：存储和高效检索这些向量
查询接口：处理用户请求并返回结果

我们选择ResNet50作为特征提取模型，它已经在ImageNet数据集上预训练，能够捕捉图像的语义特征。对于向量数据库，Milvus是当前最流行的开源选择，专为向量相似度搜索优化。

为什么选择Milvus？

支持多种相似度度量方式（余弦、欧式距离等）
提供高效的索引构建和查询算法
可扩展性强，支持分布式部署
有成熟的Python客户端

2. 环境准备与依赖安装

开始编码前，我们需要设置开发环境。建议使用Python 3.7+和最新版的Milvus（2.x版本）。

# 创建并激活虚拟环境 python -m venv img_search source img_search/bin/activate # Linux/Mac img_search\Scripts\activate # Windows # 安装核心依赖 pip install pymilvus torch torchvision pillow numpy

对于特征提取，我们将使用PyTorch提供的预训练ResNet50模型：

import torch import torchvision.models as models from torchvision import transforms # 加载预训练模型（不包含最后的全连接层） model = models.resnet50(pretrained=True) model = torch.nn.Sequential(*(list(model.children())[:-1])) model.eval() # 设置为评估模式

3. 图像特征提取流程

将图片转换为特征向量是整个系统的第一步。我们需要设计一个标准化的处理流程：

图像加载与预处理
通过神经网络提取特征
特征向量归一化

from PIL import Image def extract_features(image_path): # 定义图像预处理流程 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] ) ]) # 加载并预处理图像 img = Image.open(image_path) img_tensor = preprocess(img) img_tensor = img_tensor.unsqueeze(0) # 添加batch维度 # 提取特征 with torch.no_grad(): features = model(img_tensor) # 展平并归一化特征向量 features = features.squeeze().numpy() features = features / np.linalg.norm(features) return features

提示：特征归一化是关键步骤，能确保后续的相似度计算更加准确。归一化后的向量在进行内积运算时，结果等同于余弦相似度。

4. Milvus数据库配置与操作

现在我们来设置Milvus并创建用于存储图像向量的集合（collection）。

4.1 连接Milvus服务

from pymilvus import connections, utility # 连接到Milvus服务器 connections.connect( alias="default", host="localhost", port="19530" ) # 检查连接是否成功 if utility.has_collection("image_vectors"): utility.drop_collection("image_vectors")

4.2 创建向量集合

我们需要定义集合的schema，包括向量维度和索引类型：

from pymilvus import FieldSchema, CollectionSchema, DataType, Collection # 定义字段 fields = [ FieldSchema(name="id", dtype=DataType.INT64, is_primary=True, auto_id=True), FieldSchema(name="image_path", dtype=DataType.VARCHAR, max_length=256), FieldSchema(name="vector", dtype=DataType.FLOAT_VECTOR, dim=2048) # ResNet50输出2048维向量 ] # 创建集合schema schema = CollectionSchema( fields=fields, description="Image similarity search collection" ) # 创建集合 collection = Collection( name="image_vectors", schema=schema, using="default" )

4.3 创建高效查询索引

为了加速相似度搜索，我们需要为向量字段创建索引：

index_params = { "index_type": "IVF_FLAT", "metric_type": "IP", # 内积（等同于余弦相似度，因为向量已归一化） "params": {"nlist": 128} } collection.create_index( field_name="vector", index_params=index_params ) # 加载集合到内存 collection.load()

5. 构建完整图像检索系统

现在我们将各个组件整合成一个完整的系统。系统需要实现两个主要功能：

添加新图片到数据库
根据查询图片找出相似图片

5.1 图片入库流程

def add_image_to_db(image_path): # 提取特征向量 vector = extract_features(image_path) # 准备插入数据 data = [ [image_path], # image_path字段 [vector] # vector字段 ] # 插入数据 mr = collection.insert(data) # 刷新使数据可搜索 collection.flush() return mr.primary_keys[0]

5.2 相似图片搜索实现

搜索功能需要接收查询图片，返回最相似的若干结果：

def search_similar_images(query_image_path, top_k=5): # 提取查询图片特征 query_vector = extract_features(query_image_path) # 定义搜索参数 search_params = { "metric_type": "IP", "params": {"nprobe": 16} } # 执行搜索 results = collection.search( data=[query_vector], anns_field="vector", param=search_params, limit=top_k, output_fields=["image_path"] ) # 整理并返回结果 ret = [] for hits in results: for hit in hits: ret.append({ "image_path": hit.entity.get("image_path"), "score": hit.score }) return ret

6. 系统优化与扩展建议

基础系统搭建完成后，我们可以考虑以下优化方向：

6.1 性能优化技巧

批量插入：当需要添加大量图片时，使用批量插入显著提高效率

def batch_add_images(image_paths): vectors = [extract_features(path) for path in image_paths] data = [image_paths, vectors] mr = collection.insert(data) collection.flush() return mr.primary_keys

索引优化：根据数据量调整索引参数
- 小数据集（<1万）：IVF_FLAT
- 中等数据（1万-100万）：IVF_SQ8
- 大数据集（>100万）：HNSW

6.2 功能扩展思路

混合搜索：结合传统标签和向量相似度
实时更新：定期增量更新特征库
结果过滤：基于元数据（如时间、类别）筛选结果

6.3 部署建议

组件	推荐配置	说明
Milvus	独立服务器或Docker容器	生产环境建议分布式部署
特征提取服务	GPU服务器	使用ONNX或TensorRT加速推理
Web接口	FastAPI或Flask	提供RESTful API给前端调用

在实际项目中，我们通常会遇到各种边界情况。比如处理不同尺寸和比例的图片时，简单的中心裁剪可能丢失重要信息。一个实用的技巧是结合多种裁剪方式提取特征，然后综合结果。