• 搭建自己的以图搜图系统 (一):10 行代码以图搜图


    目前市面上有很多以图搜图的服务,如 Google 识图,百度图片搜索,淘宝拍立淘等。本文将介绍如何快速搭建自己的图片搜索引擎,只要 10 行 Python 代码就能轻松搞定!

    1. import towhee
    2. towhee.read_csv('reverse_image_search.csv') \
    3.     .runas_op['id''id'](func=lambda x: int(x)) \
    4.     .image_decode['path''img']() \
    5.     .image_embedding.timm['img''vec'](model_name='resnet50') \
    6.     .to_milvus['id''vec'](collection=collection, batch=100)
    7. towhee.glob['path']('./test/*/*.JPEG') \
    8.     .image_decode['path''img']() \
    9.     .image_embedding.timm['img''vec'](model_name='resnet50') \
    10.     .milvus_search['vec''result'](collection=collection, limit=10)

    图 1 为查询图像时的展示结果,左侧为搜索的图片,右侧是得到的相似图片。

    图 1 查询图像的展示结果

    在介绍如何搭建系统之前,我们先简单了解下以图搜图的基本原理。

    以图搜图

    以图搜图,顾名思义就是用图片搜索图片。它的应用场景包括查找原始图片、搜索相似的图片、商品搜索和推荐(根据图片在商品库中搜索同款,或者推荐相似商品)等。

    传统图像搜索¹是利用增加元数据的方法,例如:根据图片添加字幕、关键词或是文段说明,通过打标签的方式来完成检索。而以图搜图是一种基于内容的图像检索 (CBIR) 技术²,它的特点是无需关键字就能理解图像的相关内容,主要依赖于 AI 算法,目前一些排名较好的图像分类算法可以到达 99% 准确率(TOP5)³。本文将利用 AI 模型提取图像特征向量,通过特征向量计算来完成以图搜图。

    准备工作

    完整的代码已上传到 Github,欢迎读者参考和使用:build_image_search_engine.ipynb

    数据准备

    ImageNet 数据集是深度学习领域中图像分类和检测最常用数据集之一,本文的图像数据就是从中随机抽取的。我们所使用的数据集中包括训练集(train)和测试集(test)两部分,训练集有 100 个分类,每个分类有 10 张图片,测试集则是 100 个分类,每个分类 1 张图片。

    我们先下载数据并解压:

    1. curl -L https://github.com/towhee-io/examples/releases/download/data/reverse_image_search.zip -O
    2. unzip -q -o reverse_image_search.zip

    解压后会发现有一个 CSV 格式的文件,它包含了训练集中 1000 张图片的基础信息,如图像的 id、所在路径、以及类别。

    让我们以表格方式查看下文件的内容(图 2 所示):

    1. import pandas as pd
    2. df = pd.read_csv('reverse_image_search.csv')
    3. df.head()

    图 2 reverse_image_search.csv 文件的部分内容

    为了记录数据集中每个id对应的图片路径,接下来我们将读取的df转换为id_img字典。同时定义read_images函数,该函数根据搜索结果的id返回图片列表,便于最终的图片展示。

    1. import cv2
    2. from towhee._types.image import Image
    3. id_img = df.set_index('id')['path'].to_dict()
    4.     
    5. def read_images(results):
    6.     imgs = []
    7.     for re in results:
    8.         path = id_img[re.id]
    9.         imgs.append(Image(cv2.imread(path), 'BGR'))
    10.     return imgs

    至此我们完成了数据准备过程,接下来是关于图像处理的准备工作,需要用到两个重要组件 Towhee 和 Milvus。

    Towhee & Milvus

    图片搜索需要用特征向量来表征图像,我们通常利用 AI 模型提取特征向量,但面对业界的诸多模型我们该如何快速上手?"X2Vec, Towhee is all you need!",Towhee 提供开箱即用的 Embedding 流水线可以将任何非结构化数据(图像,视频,音频等)转为特征向量,通过 Towhee 我们运行一条流水线就能轻松得到特征向量。

    解决了如何提取特征向量的问题,接下来要解决的是向量搜索问题。

    想要快速简单的实现向量检索功能,选择使用 Milvus 是一个不错的技术方案。Milvus 是一个开源的向量数据库项目,它支持丰富的向量索引算法和向量计算方式,轻松实现对数百万、数十亿甚至数万亿向量的相似性搜索,具有高度灵活、稳定可靠以及高速查询等特点。

    完整的系统架构如图 3 所示,通过 Towhee + Milvus 就可以实现端到端的图像等非结构化数据分析。我们先使用 Towhee 完成非结构化数据的特征向量提取,然后 Milvus 负责存储并搜索向量,最终获取与查询数据最相似的结果并展示。

    图 3 Towhee 与 Milvus 处理图像等非结构化数据的系统架构

    理解了基于 Tohwhee 和 Milvus 的以图搜图架构,接下来我们要先完成 Towhee 和 Milvus 的安装:

    注意:Milvus 支持单机安装集群安装,本文使用 docker-compose 方式安装单机 Milvus,在此之前请先检查本机环境的软硬件条件

    1. #安装 Towhee
    2. pip install towhee
    3. #安装单机版 Milvus
    4. wget https://github.com/milvus-io/milvus/releases/download/v2.0.2/milvus-standalone-docker-compose.yml -O docker-compose.yml
    5. docker-compose up -d
    • Towhee

    Towhee 支持图像 Embedding,音频 Embedding,视频 Embedding 等非结构化数据特征提取的方法,这些都被称为 Towhee 的算子(Operator),算子是流水线(Pipeline)中的单个节点,一个图像特征提取流水线就可以通过连接 image_decode 算子和 image_embedding.timm 算子实现,其中 Embedding 算子可以通过指定model_name="resnet50"利用 ResNet50 模型生成特征向量(结果如图 4 所示):

    1. import towhee
    2. towhee.glob['path']('./test/lion/n02129165_13728.JPEG') \
    3.       .image_decode['path''img']() \
    4.       .image_embedding.timm['img''vec'](model_name='resnet50') \
    5.       .select['img''vec']() \
    6.       .show()

    图 4 ResNet50 特征向量提取

    image_embedding.timm 算子支持各种预训练好的模型,包括vgg16resnet50vit_base_patch8_224convnext_base等。该算子被托管在 Towhee Hub 上,Hub 上有成百上千个算子,你可以在其中找到任何你想要的 Embedding 处理方式。

    • Milvus

    接下来在 Milvus 数据库中创建集合(Collection),集合中的 Fields 包含两列:id 和 embedding,其中 id 是集合的主键。另外我们可以为 embedding 创建 IVF_FLAT 基于量化的索引,其中索引的参数是 nlist=2048,计算方式是 "L2" 欧式距离:

    1. from pymilvus import connections, FieldSchema, CollectionSchema, DataType, Collection, utility
    2. def create_milvus_collection(collection_name, dim):
    3.     connections.connect(host='127.0.0.1', port='19530')
    4.     
    5.     if utility.has_collection(collection_name):
    6.         utility.drop_collection(collection_name)
    7.     
    8.     fields = [
    9.     FieldSchema(name='id', dtype=DataType.INT64, descrition='ids', is_primary=True, auto_id=False),
    10.     FieldSchema(name='embedding', dtype=DataType.FLOAT_VECTOR, descrition='embedding vectors', dim=dim)
    11.     ]
    12.     schema = CollectionSchema(fields=fields, description='reverse image search')
    13.     collection = Collection(name=collection_name, schema=schema)
    14.     # create IVF_FLAT index for collection.
    15.     index_params = {
    16.         'metric_type':'L2',
    17.         'index_type':"IVF_FLAT",
    18.         'params':{"nlist":2048}
    19.     }
    20.     collection.create_index(field_name="embedding", index_params=index_params)
    21.     return collection
    22. collection = create_milvus_collection('reverse_image_search'2048)

    运行以图搜图并展示

    如图 5 所示,我们将以图搜图服务分为插入和查询两步:首先在 Towhee hub 上选择所需的图像 Embedding 流水线,用来提取图像数据集的特征向量,再将特征向量存入 Milvus 中;查询的时候利用同样的流水线提取查询图像的特征向量,然后在 Milvus 中检索得出相似的结果,最终展示出图片。

    图5 以图搜图流程图

    图像数据入库

    Towhee 不光拥有丰富的算子来处理非结构化数据,还提供了简单好用的接口来处理各种数据,当然也集成了 Milvus 的一些基本用法,通过在“流水线”中连接这些算子或接口,图像入库操作将变得十分简单。

    图像数据入库”流水线“的代码说明参考注释,如果你有不明白的地方,欢迎留言讨论,或者直接给 Towhee 项目提 ISSUE。

    1. import towhee 
    2.  
    3. dc = ( 
    4.     towhee.read_csv('reverse_image_search.csv'#读取 CSV 格式的表格,包含了 id,path 和 label 列
    5.  .runas_op['id''id'](func=lambda x: int(x)) #将每一行的 id 从 str 类型转为 int 类型
    6.  .image_decode['path''img']() #读取每一行 path 对应的图像,并将其解码为 Towhee 的图像格式 
    7.  .image_embedding.timm['img''vec'](model_name='resnet50'#提取特征向量
    8.  .tensor_normalize['vec''vec']() #将向量进行归一化
    9.  .to_milvus['id''vec'](collection=collection, batch=100#将 id 和 vec 批量 100 条插入到 Milvus 集合
    10. )

    查询图像并展示

    查询图像时需要的图像处理算子与前面类似,包括image_decodeimage_embedding.timmtensor_normalize,而在最后分析检索结果时,需用到数据准备部分定义好的read_images函数,通过指定runas_op中的func将该函数加入到 Towhee 流水线中。

    查询图像”流水线“的代码说明参考注释,如果你有不明白的地方,欢迎留言讨论,或者直接给 Towhee 项目提 ISSUE。

    1. (towhee.glob['path']('./test/w*/*.JPEG'#读取满足指定模式下的所有图片数据为 path 
    2.  .image_decode['path''img']() #读取每一行 path 对应的图像,并将其解码为 Towhee 的图像格式  
    3.  .image_embedding.timm['img''vec'](model_name='resnet50'#提取特征向量
    4.  .tensor_normalize['vec''vec']() #将向量进行归一化
    5.  .milvus_search['vec''result'](collection=collection, limit=5#在 Milvus 集合中搜索向量,并返回结果
    6.  .runas_op['result''result_img'](func=read_images) #处理 Milvus 的检索结果,最终返回图像用于展示
    7.  .select['img''result_img']() #选择指定列; 
    8.  .show()
    9. )

    当代码执行完毕之后,我们将得到类似下面的结果。

    图 6 查询图像的展示结果

    Gradio 部署服务

    Gradio 为机器学习模型提供 Web 演示界面,我们所熟知的 Huggingface 的 Demo 界面也是利用它实现的。Gradio 支持上传和展示图片,如图 7 所示,我们可以利用它实现一个以图搜图可交互的服务,可以在它生成的 Web 界面中上传一张图片,来查询与其相似的其他图片。

     图 7 基于 Gradio 的以图搜图服务.

    Gradio 的使用十分简单,我们只需要定义查询图片的函数image_search_function,确定其输入和输出,在创建 Gradio 服务的时候绑定函数与对应的输入输出,最后启动服务就搞定了!

    1. import gradio
    2. from towhee.types.image_utils import from_pil
    3. with towhee.api() as api:
    4.     image_search_function = (
    5.         api.runas_op(func=lambda img: from_pil(img))
    6.             .image_embedding.timm(model_name='resnet50')
    7.             .tensor_normalize()
    8.             .milvus_search(collection='reverse_image_search', limit=5)
    9.             .runas_op(func=lambda res: [id_img[x.idfor x in res])
    10.             .as_function()
    11.     )
    12.     
    13. interface = gradio.Interface(image_search_function, 
    14.                              gradio.inputs.Image(type="pil", source='upload'),
    15.                              [gradio.outputs.Image(type="file", label=Nonefor _ in range(5)]
    16.                             )
    17. interface.launch(inline=True)

    当你成功启动了 Gradio,前端页面会嵌入到当前 notebook 中,如图 7 所示,你也可以点击 Gradio 提供的链接(http://127.0.0.1:7860)打开前端,用于查询图片并查看结果。

    如果你想要把以图搜图服务和朋友分享,也可以尝试在launch函数中设置参数 share=True,这时会得到一个公共网址,诸如https://xxxx.gradio.app,我们就可以把自己搭建的以图搜图服务分享给小伙伴啦。

    写在最后

    其实 Towhee 不仅仅能处理图片这种非结构化数据,对于音频、视频等数据也能进行分析处理,参考本文的实现,同样我们也可以 10 行代码来实现音频处理、视频处理等相关业务的 AI 服务,感兴趣的话大家可以自行尝试。

    在下一篇内容中,我将介绍如何对这个系统进行调优处理,敬请期待。

    参考

    [1] https://zh.wikipedia.org/wiki/%E5%9C%96%E5%83%8F%E6%AA%A2%E7%B4%A2

    [2] https://en.wikipedia.org/wiki/Reverse_image_search

    [3] https://paperswithcode.com/sota/image-classification-on-imagenet?metric=Top%205%20Accuracy


    更多项目更新及详细内容请关注我们的项目( https://github.com/towhee-io/towhee/blob/main/towhee/models/README_CN.md) ,您的关注是我们用爱发电的强大动力,欢迎 star, fork, slack 三连 :)

    zilliz用户交流

     

  • 相关阅读:
    YAMLException : java.nio.charset.MalformedInputException : Input length = 1
    笔记/日记应用 memos
    笔记本电脑充电器、电源适配器以及Type-c手机充电器UL60950认证报告具体要求
    区块链工作原理(节点、层)
    新能源汽车发展迅猛,市场份额已突破50%
    算法 合并有序序列
    织信Informat如何连接其他应用?
    nginx - 负载均衡配置-负载均衡策略
    [超详细] GraalVM打包含有JNI的本地镜像
    [MySQL]单行函数
  • 原文地址:https://blog.csdn.net/weixin_44839084/article/details/126030438