简介:图像检索系统是一份面向计算机视觉初学者与课程实践者的完整源码资源,围绕以图搜图这一典型任务,覆盖图像预处理、特征提取、相似度计算、索引构建、检索策略与性能评估等核心环节,可用于搜索引擎、数字图书馆、医学影像分析等场景的学习与二次开发。压缩包共48个文件,约1.18MB,以31个bmp样本图片和5个h头文件、3个cpp源文件为主体,另含工程解决方案、资源脚本与可执行程序,构成可直接编译运行的工程结构,便于对照代码理解各模块实现。已有1209人学习下载,说明其在教学与练手场景中具有一定参考价值。读者可通过阅读与修改源码,掌握颜色直方图、局部描述符等特征的组织方式,理解欧氏距离、余弦相似度等度量方法,并借助索引结构体会检索效率的优化思路,从而形成从特征到检索的完整认知。
1. 图像检索系统:从“以图搜图”到可复现的工程落地
上周有个做电商的朋友甩过来一个需求:用户上传一张鞋子的照片,要在几万张商品图里把同款或相似款捞出来,响应时间还得压在 200ms 以内。他第一反应是“调个大模型 API 不就完了”,结果一算账,几万张图每次全量比对,成本和延迟都扛不住。这就是图像检索系统要解决的核心问题——不是“能不能搜”,而是“怎么在可控成本下搜得准、搜得快”。这份资源包给的正是一套从特征提取、向量索引到检索服务的完整链路,适合两类人:一是想搞懂以图搜图底层逻辑的开发者,二是手里有图库、需要快速搭一套可用检索服务的从业者。它不追求 SOTA 精度,但胜在链路完整、参数可调、能跑通。
2. 特征提取与向量化:检索系统的地基怎么打
图像检索系统的本质,是把“看图”这件事翻译成“算距离”。一张图进来,先过一个特征提取器,输出一个固定维度的向量,然后拿这个向量去向量库里找最近的邻居。所以整条链路里,特征提取器选什么、向量怎么归一化、维度定多少,直接决定了后面检索的天花板。这一章先把地基讲清楚,再落到具体代码。
2.1 特征提取器的选型:CNN 还是 ViT,预训练权重怎么挑
常见做法是分两档。第一档是轻量场景,用 ResNet50 这类 CNN 骨干,去掉最后的全连接分类层,取全局平均池化后的 2048 维输出作为图像特征。它的好处是推理快、显存占用低,单张图在普通 GPU 上几毫秒就能出向量,适合图库规模在十万级以内、对精度要求不是极致的场景。第二档是精度优先,用 ViT-B/16 或者 CLIP 的图像编码器,输出 512 维或 768 维向量,语义表达能力明显更强,尤其是跨模态检索(用文字搜图)时优势很大,代价是推理慢一些、对显存更敏感。
我一般会先问一句:你的检索是“同款精确匹配”还是“相似风格召回”?前者用 ResNet 系列足够,后者建议上 ViT 或 CLIP。这里有个容易被忽略的点——预训练权重的来源要和你的图库领域尽量接近。用 ImageNet 权重去提商品图特征,效果通常还行;但如果你的图库是医学影像、工业缺陷图这类分布差异大的数据,最好先用少量标注数据做一轮微调,哪怕只微调最后两三个 block,召回率也能有明显提升。
下面这段是特征提取的核心代码,用 PyTorch 写,ResNet50 和 ViT 两种骨干都留了切换口:
import torch import torch.nn as nn from torchvision import models, transforms from PIL import Image class FeatureExtractor: def __init__(self, backbone="resnet50", device="cuda"): self.device = device if backbone == "resnet50": model = models.resnet50(weights=models.ResNet50_Weights.IMAGENET1K_V2) # 去掉最后的全连接层,保留全局平均池化输出 self.model = nn.Sequential(*list(model.children())[:-1]) self.dim = 2048 elif backbone == "vit_b16": model = models.vit_b_16(weights=models.ViT_B_16_Weights.IMAGENET1K_V1) # ViT 取 CLS token 对应的输出作为图像特征 self.model = model self.dim = 768 self.model.eval().to(device) # 预处理必须和预训练权重训练时保持一致 self.tf = 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]), ]) @torch.no_grad() def extract(self, img_path): img = Image.open(img_path).convert("RGB") x = self.tf(img).unsqueeze(0).to(self.device) if self.dim == 768: # ViT 分支 feat = self.model(x) else: feat = self.model(x).flatten(1) # L2 归一化,后面用内积等价于余弦相似度 feat = torch.nn.functional.normalize(feat, p=2, dim=1) return feat.cpu().numpy().astype("float32")这段代码有三个参数值得说清楚。backbone决定维度和速度,ResNet50 出 2048 维、ViT-B/16 出 768 维,维度越高索引内存越大,十万张图 2048 维 float32 大约占 800MB,768 维只要 300MB 左右。Normalize里的均值和方差是 ImageNet 统计值,换了预训练权重就得跟着换,这一步错了特征会整体偏移,检索结果会莫名其妙地差。最后的 L2 归一化是关键,归一化之后向量内积就等于余弦相似度,索引库那边可以直接用内积检索,省掉一次距离换算。
2.2 批量向量化与落盘:把图库变成特征矩阵
单张提取跑通之后,下一步是把整个图库批量转成特征矩阵,并且把图片路径和向量一一对应地存下来。这一步的坑在于内存和 IO:几万张图如果一张张读、一张张存,速度慢不说,还容易因为某张图损坏导致整个任务中断。我一般会做三件事——批量读、异常跳过、分片落盘。
import os import numpy as np from tqdm import tqdm def build_feature_bank(img_dir, extractor, batch_size=32, shard_size=10000): paths = [] for root, _, files in os.walk(img_dir): for f in files: if f.lower().endswith((".jpg", ".jpeg", ".png")): paths.append(os.path.join(root, f)) paths.sort() all_feats, all_paths = [], [] for i in tqdm(range(0, len(paths), batch_size)): batch = paths[i:i + batch_size] for p in batch: try: feat = extractor.extract(p) all_feats.append(feat) all_paths.append(p) except Exception as e: # 损坏图直接跳过并记录,不要让整个任务挂掉 print(f"[skip] {p}: {e}") # 每攒够一个分片就落盘,避免内存爆掉 if len(all_feats) >= shard_size: _dump(all_feats, all_paths, len(paths)) all_feats, all_paths = [], [] if all_feats: _dump(all_feats, all_paths, len(paths)) def _dump(feats, paths, total): mat = np.vstack(feats) idx = len(os.listdir("./feature_bank")) if os.path.exists("./feature_bank") else 0 os.makedirs("./feature_bank", exist_ok=True) np.save(f"./feature_bank/feat_{idx}.npy", mat) with open(f"./feature_bank/paths_{idx}.txt", "w") as f: f.write("\n".join(paths)) print(f"shard {idx} saved: {mat.shape}, total={total}")batch_size控制的是每轮读图的批量,GPU 显存小就调到 16 甚至 8;shard_size控制的是每多少条落一次盘,图库上万时建议设成 10000 左右,这样单个 npy 文件不会太大,后续建索引也方便分片加载。异常跳过那段是血泪经验——图库里总会有几张下载不完整或者格式损坏的图,不处理的话跑到一半崩掉,前面几万张白跑。落盘时特征矩阵和路径文件分开存,靠行号对齐,这是后面建索引和回查原图的基础,千万别把顺序搞乱。
3. 向量索引与相似度检索:Faiss 怎么配才不翻车
特征矩阵有了,接下来要解决“给定一个查询向量,怎么在几万甚至几百万条里快速找到最近的 K 条”。暴力比对当然可以,但图库一大就顶不住。这一章讲索引选型、参数怎么调,以及检索结果怎么回查原图。
3.1 索引类型选择:Flat、IVF、HNSW 的取舍
Faiss 里常用的索引就那几种,选错了要么慢要么不准。IndexFlatIP是暴力内积检索,精度 100%,十万级以内完全够用,建索引就是一次矩阵拷贝,没有训练过程,最省心。IndexIVFFlat是倒排索引,先把向量聚成若干簇,查询时只扫最近的几个簇,速度快但会损失一点召回,需要先用样本训练聚类中心。IndexHNSWFlat是图索引,查询快、召回高,但建索引慢、内存占用大。
我的经验是:图库五万以内直接上 Flat,别折腾;五万到五十万用 IVF,nlist设成sqrt(N)量级,查询时nprobe从 8 开始往上调,调到召回率满足要求为止;五十万以上再考虑 HNSW,但要注意它吃内存,2048 维的向量建 HNSW 内存会翻好几倍。下面这段是建索引和查询的完整示例:
import faiss import numpy as np def build_index(feat_path, index_type="flat", nlist=256): mat = np.load(feat_path).astype("float32") dim = mat.shape[1] if index_type == "flat": index = faiss.IndexFlatIP(dim) # 向量已归一化,内积即余弦 index.add(mat) elif index_type == "ivf": quantizer = faiss.IndexFlatIP(dim) index = faiss.IndexIVFFlat(quantizer, dim, nlist, faiss.METRIC_INNER_PRODUCT) index.train(mat) # IVF 必须先训练 index.add(mat) index.nprobe = 16 # 查询时扫描的簇数 return index def search(index, query_feat, topk=10): query_feat = query_feat.astype("float32") scores, ids = index.search(query_feat, topk) return scores[0], ids[0]nlist是聚类簇数,经验值是4*sqrt(N)到16*sqrt(N)之间,十万条图取 256 到 1024 都合理,太小了每个簇里向量多、查询慢,太大了训练时间长、召回反而可能掉。nprobe是查询时扫几个簇,它和召回率、延迟是直接 trade-off,线上一般从 8 开始压测,逐步加到延迟逼近预算为止。还有个细节:IVF 的train必须用足够多的样本,样本数至少是nlist的几十倍,否则聚类中心训不准,检索结果会飘。
3.2 检索结果回查与重排:别让 top-1 骗了你
向量检索返回的是 id 和相似度分数,id 对应特征矩阵的行号,要靠路径文件映射回原图。这一步本身不难,但有个常见误区——很多人拿到 top-10 就直接展示,结果用户发现排第一的并不是最像的。原因是特征提取器对某些细粒度差异不敏感,比如同款不同色、同色不同材质。我的做法是加一层轻量重排:先用向量召回 top-50,再用一个更细的特征(比如局部区域特征或者颜色直方图)对这 50 条重新打分,取前 10 展示。
def rerank(query_path, candidate_paths, extractor, topk=10): q_feat = extractor.extract(query_path) scores = [] for p in candidate_paths: c_feat = extractor.extract(p) # 余弦相似度,向量已归一化,直接内积 sim = float(np.dot(q_feat, c_feat.T)) scores.append((p, sim)) scores.sort(key=lambda x: x[1], reverse=True) return scores[:topk]重排的代价是候选集里每张图都要重新提一次特征,50 条的话延迟增加几十毫秒,换来的是 top-10 的排序质量明显提升。如果延迟预算紧,可以把重排特征换成更轻的(比如只算颜色直方图距离),或者把重排候选缩到 top-20。这里提醒一句:重排用的特征和召回用的特征最好不是同一个,否则重排等于没排,只是把同样的分数又算了一遍。
4. 服务化与性能调优:从脚本到能扛请求的接口
前面两章跑通的是离线链路,真正上线还得包成服务,处理并发请求、控制延迟、管理索引加载。这一章讲怎么用 FastAPI 把检索能力暴露成接口,以及几个影响吞吐的关键参数。
4.1 FastAPI 封装检索接口:加载一次,常驻内存
索引和特征提取模型都是重资源,绝对不能每个请求都重新加载。正确做法是服务启动时加载一次,挂在全局或者应用状态里,请求进来直接复用。下面是一个最小可用的接口实现:
from fastapi import FastAPI, UploadFile, File from contextlib import asynccontextmanager import numpy as np import faiss state = {} @asynccontextmanager async def lifespan(app: FastAPI): # 启动时加载模型和索引,只做一次 state["extractor"] = FeatureExtractor(backbone="resnet50") state["index"] = faiss.read_index("./index/faiss.index") state["paths"] = open("./feature_bank/paths_0.txt").read().splitlines() yield state.clear() app = FastAPI(lifespan=lifespan) @app.post("/search") async def search(file: UploadFile = File(...), topk: int = 10): tmp = f"/tmp/{file.filename}" with open(tmp, "wb") as f: f.write(await file.read()) feat = state["extractor"].extract(tmp) scores, ids = state["index"].search(feat, topk) results = [{"path": state["paths"][i], "score": float(s)} for i, s in zip(ids[0], scores[0]) if i != -1] return {"results": results}lifespan是 FastAPI 新版推荐的启动/关闭钩子,比老式的@app.on_event更干净。topk做成查询参数,方便前端按需调整。注意ids里可能出现-1,那是 Faiss 在候选不足时填充的无效 id,回查路径前必须过滤掉,否则会数组越界——这个坑我见过不止一次。另外上传文件先落临时目录再读,是因为特征提取器吃的是文件路径,如果你的提取器改成吃 PIL Image 对象,可以省掉这次磁盘 IO。
4.2 延迟与吞吐的调优参数:并发、批处理、索引预热
接口跑通只是第一步,压测时才会暴露真问题。几个关键旋钮:一是uvicorn的 worker 数,CPU 核数的一半到核数之间比较稳,开太多反而因为上下文切换掉吞吐;二是特征提取的批处理,如果前端支持一次传多张图,把 batch 拼起来过一次模型比逐张快很多;三是索引预热,服务刚启动时第一次查询会明显偏慢,因为索引数据还没进页缓存,可以在启动后主动跑几条假查询把热路径走一遍。
提示:Faiss 索引文件用
read_index加载后是只读的,如果图库有增量更新,不要试图原地修改,正确做法是定期重建索引然后原子替换文件,服务端监听文件变更后重新加载。
延迟预算怎么拆:特征提取通常占大头,ResNet50 单张在 GPU 上 5 到 10ms,CPU 上可能到 50ms 以上;索引查询 Flat 十万级在 10ms 内,IVF 更快;网络和序列化几毫秒。所以如果延迟要求严,优先保证特征提取跑在 GPU 上,其次才是调索引参数。吞吐方面,单卡 GPU 跑 ResNet50 大概能到每秒几百张,够大多数中小图库用了。
5. 避坑与排查:那些让检索结果“玄学”的细节
链路跑通不代表结果可信,下面这几条是我在实际项目里反复踩过的,按“现象 → 原因 → 解决”列出来,照着排查能省不少时间。
5.1 召回结果整体偏移,相似图排不上去
现象是查询一张明显有同款的图,top-10 里却混进一堆不相干的。原因大概率出在预处理不一致——建库时用的归一化均值方差和查询时用的不一样,或者建库用了CenterCrop查询用了Resize,特征空间直接错位。解决办法是把预处理逻辑抽成一个函数,建库和查询共用同一份代码,别两处各写一遍。排查时可以拿建库时的某张图当查询图,如果它自己排不到 top-1,基本就是预处理或归一化出了问题。
5.2 IVF 索引召回率忽高忽低
现象是同样的查询,有时结果正常有时明显变差。原因是nprobe设得太小,查询向量落在簇边界附近时,真正相似的向量在相邻簇里没被扫到。解决是把nprobe调大,或者改用IndexIVFPQ之外的IndexIVFFlat保证簇内精度。另一个隐蔽原因是训练样本太少,聚类中心没训好,这种情况只能补样本重训,调参救不回来。
5.3 服务跑一段时间内存持续上涨
现象是接口刚启动正常,跑几小时后内存越吃越多,最后 OOM。原因通常是请求里创建的临时文件没清理,或者特征提取时张量没释放。解决是临时文件用完即删,提取函数里加torch.no_grad()并且显式del掉中间张量,必要时定期调torch.cuda.empty_cache()。还有一种情况是索引被反复加载,检查一下是不是每次请求都触发了重载逻辑。
5.4 增量更新后新旧图检索结果不一致
现象是往图库加了一批新图,重建索引后老图的排序变了。原因是特征矩阵的行号变了,路径文件和特征矩阵没对齐,或者重建时用了不同的分片顺序。解决是重建索引时严格按路径排序后重新生成特征矩阵,路径文件和 npy 文件同批次落盘,别拿旧路径文件配新特征矩阵。
5.5 跨域图片检索效果差
现象是商品图库用得好好的,换成用户上传的实拍图就搜不准。原因是训练域和查询域分布差异大,模型对背景、光照、拍摄角度不鲁棒。解决有两个方向:一是查询侧做预处理,裁掉明显背景、做颜色归一化;二是用少量目标域数据微调特征提取器,哪怕只微调最后几层,域适应效果也很明显。
6. 进阶技巧:用查询扩展和索引分片把召回再拉一档
基础链路跑顺之后,如果还想在召回率上再挤一点,有两个成本不高但见效明显的技巧。第一个是查询扩展:用户上传一张图,先做水平翻转、轻微裁剪、亮度扰动,生成三到五个变体,每个变体各查一次,把结果按 id 合并去重、分数取最大。这相当于用数据增强对抗拍摄角度和光照差异,实测在实拍图场景下 top-10 召回能提升几个百分点。代价是查询次数翻几倍,适合离线评估或者对延迟不敏感的批量任务。
def query_expansion(img_path, extractor, index, topk=10): from PIL import Image, ImageEnhance variants = [img_path] img = Image.open(img_path).convert("RGB") # 水平翻转 img.transpose(Image.FLIP_LEFT_RIGHT).save("/tmp/flip.jpg") variants.append("/tmp/flip.jpg") # 亮度扰动 ImageEnhance.Brightness(img).enhance(1.2).save("/tmp/bright.jpg") variants.append("/tmp/bright.jpg") merged = {} for v in variants: feat = extractor.extract(v) scores, ids = index.search(feat, topk) for i, s in zip(ids[0], scores[0]): if i == -1: continue merged[i] = max(merged.get(i, -1), float(s)) ranked = sorted(merged.items(), key=lambda x: x[1], reverse=True) return ranked[:topk]第二个技巧是索引分片。图库到百万级以后,单文件索引加载慢、内存压力大,可以按类别或者按哈希把图库切成若干片,每片一个索引文件,查询时并行查所有分片再归并。分片的好处是增量更新只动一个片,重建成本低;坏处是查询要扫多个索引,归并逻辑得自己写。分片数一般取 4 到 16,太少没意义,太多归并开销盖过收益。
| 技巧 | 适用场景 | 召回提升 | 额外成本 |
|---|---|---|---|
| 查询扩展 | 实拍图、跨域检索 | 中等 | 查询次数翻倍 |
| 索引分片 | 百万级以上图库 | 无(保精度) | 归并逻辑、多文件管理 |
| 重排 | 细粒度区分 | 较高 | 候选集重新提特征 |
这两个技巧我一般不会一上来就上,先把基础链路的预处理一致性和索引参数调稳,确认瓶颈在召回而不是工程问题之后,再考虑加。有次我图省事直接上了查询扩展,结果发现真正的问题是建库时归一化写错了,扩展只是把错误放大。从那以后我每次改检索链路,都强制先拿建库图自查询验证 top-1 是不是自己,再谈优化。希望这些细节能帮你在搭自己的图像检索系统时少走点弯路。
本文还有配套的精品资源,点击获取