2、图片对比及图片向量提取(代码)

48 阅读2分钟

以下代码将演示如何提取图片的特征向量,可配合向量数据库实现图片的搜索功能。

github下载代码:参考image/image_deduplication目录下的图片文件github.com/yun81911/ex…

以下是我导入数据并进行查询问答查询的代码:

# 导入 towhee 库中的 pipe 和 ops 模块,用于构建数据处理管道和操作
from towhee import pipe, ops
# 导入 towhee 库中的 DataCollection 类,用于封装和展示管道执行结果
from towhee.datacollection import DataCollection

# 定义一个名为 emb_pipe 的数据处理管道,用于图像特征嵌入
emb_pipe = (
    # 定义管道的输入,包含两个参数:'target' 表示目标图像,'path' 表示待比较图像的路径列表
    pipe.input('target', 'path')
    # 对 'path' 进行扁平化处理,将可能的嵌套列表展开为一维列表
    .flat_map('path', 'path', lambda x: x)
    # 对 'target' 进行图像解码操作,使用 OpenCV 的 cv2 库,将图像解码为 RGB 格式
    .map('target', 'target', ops.image_decode.cv2('rgb'))
    # 对 'path' 进行图像解码操作,使用 OpenCV 的 cv2 库,将图像解码为 RGB 格式
    .map('path', 'image', ops.image_decode.cv2('rgb'))
    # 对解码后的 'target' 图像进行特征嵌入操作,使用 timm 库中的 resnet50 模型,将图像转换为特征向量
    .map('target', 'target_emb', ops.image_embedding.timm(model_name='resnet50'))
    # 对解码后的 'image' 图像进行特征嵌入操作,使用 timm 库中的 resnet50 模型,将图像转换为特征向量
    .map('image', 'emb', ops.image_embedding.timm(model_name='resnet50'))
)

# 执行管道,传入输入参数 'Lenna.png' 作为目标图像,['Lenna.png', 'logo.png'] 作为待比较图像的路径列表
# 并指定输出的字段,包括目标图像、待比较图像、图像路径、待比较图像的特征向量、目标图像的特征向量
emb_res = emb_pipe.output('target', 'image', 'path', 'emb', 'target_emb')('Lenna.png', ['Lenna.png', 'logo.png'])

# 使用 DataCollection 类对管道执行结果进行封装,并调用 show() 方法展示结果
emb_res_collection = DataCollection(emb_res)
emb_res_collection.show()
# 打印 Lenna 及 logo 的特征向量值
for result in emb_res_collection:
    path = result.path
    emb = result.emb
    print(f"Image path: {path}, Feature vector value: {emb}")


# 打印 Lenna 及 logo 的图片比较结果
import numpy as np
thresh = 0.01
detect_res = (
    emb_pipe.map(('emb', 'target_emb'), 'is_similar', lambda x, y: np.linalg.norm(x - y) < thresh)
        .output('target', 'image', 'is_similar')
)

res = detect_res('Lenna.png', ['Lenna.png', 'logo.png'])
DataCollection(res).show()