乐趣区

关于数据库:搭建自己的以图搜图系统-一10-行代码搞定以图搜图

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

import towhee

towhee.read_csv('reverse_image_search.csv') \
    .runas_op['id', 'id'](func=lambda x: int(x)) \
    .image_decode['path', 'img']() \
    .image_embedding.timm['img', 'vec'](model_name='resnet50') \
    .to_milvus['id', 'vec'](collection=collection, batch=100)

towhee.glob['path']('./test/*/*.JPEG') \
    .image_decode['path', 'img']() \
    .image_embedding.timm['img', 'vec'](model_name='resnet50') \
    .milvus_search['vec', 'result'](collection=collection, limit=10)

图 1 为查问图像时的展现后果,左侧为搜寻的图片,右侧是失去的类似图片。

在介绍如何搭建零碎之前,咱们先简略理解下以图搜图的基本原理。

以图搜图

以图搜图,顾名思义就是用图片搜寻图片。它的利用场景包含查找原始图片、搜寻类似的图片、商品搜寻和举荐(依据图片在商品库中搜寻同款,或者举荐类似商品)等。

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

筹备工作

残缺的代码已上传到 Github,欢送读者参考和应用:build_image_search_engine.ipynb。

数据筹备

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

咱们先下载数据并解压:

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

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

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

import pandas as pd

df = pd.read_csv('reverse_image_search.csv')
df.head()

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

import cv2
from towhee._types.image import Image

id_img = df.set_index('id')['path'].to_dict()
    
def read_images(results):
    imgs = []
    for re in results:
        path = id_img[re.id]
        imgs.append(Image(cv2.imread(path), 'BGR'))
    return imgs

至此咱们实现了数据筹备过程,接下来是对于图像处理的筹备工作,须要用到两个重要组件 Towhee 和 Milvus。

Towhee & Milvus

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

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

想要疾速简略的实现向量检索性能,抉择应用 Milvus 是一个不错的技术计划。Milvus 是一个开源的向量数据库我的项目,它反对丰盛的向量索引算法和向量计算形式,轻松实现对数百万、数十亿甚至数万亿向量的相似性搜寻,具备高度灵便、稳固牢靠以及高速查问等特点。

残缺的零碎架构如图 3 所示,通过 Towhee + Milvus 就能够实现端到端的图像等非结构化数据分析。咱们先应用 Towhee 实现非结构化数据的特征向量提取,而后 Milvus 负责存储并搜寻向量,最终获取与查问数据最类似的后果并展现。

了解了基于 Tohwhee 和 Milvus 的以图搜图架构,接下来咱们要先实现 Towhee 和 Milvus 的装置:

留神:Milvus 反对单机装置和集群装置,本文应用 docker-compose 形式装置单机 Milvus,在此之前请先查看本机环境的软硬件条件。

# 装置 Towhee
$ pip install towhee

#装置单机版 Milvus
$ wget https://github.com/milvus-io/milvus/releases/download/v2.0.2/milvus-standalone-docker-compose.yml -O docker-compose.yml
$ docker-compose up -d
  • Towhee

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

import towhee

towhee.glob['path']('./test/lion/n02129165_13728.JPEG') \
      .image_decode['path', 'img']() \
      .image_embedding.timm['img', 'vec'](model_name='resnet50') \
      .select['img', 'vec']() \
      .show()

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” 欧式间隔:

from pymilvus import connections, FieldSchema, CollectionSchema, DataType, Collection, utility

def create_milvus_collection(collection_name, dim):
    connections.connect(host='127.0.0.1', port='19530')
    
    if utility.has_collection(collection_name):
        utility.drop_collection(collection_name)
    
    fields = [FieldSchema(name='id', dtype=DataType.INT64, descrition='ids', is_primary=True, auto_id=False),
    FieldSchema(name='embedding', dtype=DataType.FLOAT_VECTOR, descrition='embedding vectors', dim=dim)
    ]
    schema = CollectionSchema(fields=fields, description='reverse image search')
    collection = Collection(name=collection_name, schema=schema)

    # create IVF_FLAT index for collection.
    index_params = {
        'metric_type':'L2',
        'index_type':"IVF_FLAT",
        'params':{"nlist":2048}
    }
    collection.create_index(field_name="embedding", index_params=index_params)
    return collection

collection = create_milvus_collection('reverse_image_search', 2048)

运行以图搜图并展现

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

图像数据入库

Towhee 不光领有丰盛的算子来解决非结构化数据,还提供了简略好用的接口来解决各种数据,当然也集成了 Milvus 的一些根本用法,通过在“流水线”中连贯这些算子或接口,图像入库操作将变得非常简略。

图像数据入库”流水线“的代码阐明参考正文,如果你有不明确的中央,欢送留言探讨,或者间接给 Towhee 我的项目提 ISSUE。

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

查问图像并展现

查问图像时须要的图像处理算子与后面相似,包含 image_decodeimage_embedding.timmtensor_normalize,而在最初剖析检索后果时,需用到数据筹备局部定义好的 read_images 函数,通过指定 runas_op 中的 func 将该函数退出到 Towhee 流水线中。

查问图像”流水线“的代码阐明参考正文,如果你有不明确的中央,欢送留言探讨,或者间接给 Towhee 我的项目提 ISSUE。

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

当代码执行结束之后,咱们将失去相似上面的后果。

Gradio 部署服务

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

Gradio 的应用非常简略,咱们只须要定义查问图片的函数image_search_function,确定其输出和输入,在创立 Gradio 服务的时候绑定函数与对应的输入输出,最初启动服务就搞定了!

import gradio
from towhee.types.image_utils import from_pil

with towhee.api() as api:
    image_search_function = (api.runas_op(func=lambda img: from_pil(img))
            .image_embedding.timm(model_name='resnet50')
            .tensor_normalize()
            .milvus_search(collection='reverse_image_search', limit=5)
            .runas_op(func=lambda res: [id_img[x.id] for x in res])
            .as_function())
    

interface = gradio.Interface(image_search_function, 
                             gradio.inputs.Image(type="pil", source='upload'),
                             [gradio.outputs.Image(type="file", label=None) for _ in range(5)]
                            )
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…

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

[3] https://paperswithcode.com/so…

更多我的项目更新及具体内容请关注咱们的我的项目(https://github.com/towhee-io/…),您的关注是咱们用爱发电的弱小能源,欢送 star, fork, slack 三连 :)

退出移动版