{模型之力} 图片相似度计算,图搜图,文搜图 By similarities
评论
收藏

{模型之力} 图片相似度计算,图搜图,文搜图 By similarities

经验分享
K2022
2024-02-28 20:56·浏览量:1578
K2022
影刀认证工程师
发布于 2024-02-28 20:561578浏览

此贴用于补全上个帖子( 文本相似度计算 By similarities )没有介绍的功能:

使用前需要 安装 (similarities ,torch ,pydantic)这些库

下面介绍的所有功能都有CPU和GPU模式(默认CPU模式,GPU模式将参数device设置成cuda即可)

注:GPU模式需要额外安装一些依赖环境,使用前需要配置好,环境配置参考我之前个帖子中GPU模式的说明: https://www.yingdao.com/community/detaildiscuss?id=e4c9c725-7c7c-4cd8-952f-1aeac275d7a9&tag=&from=userCenter&sort=createTime&page=1

代码实现:

1:图片相似度计算

# 使用提醒:
# 1. xbot包提供软件自动化、数据表格、Excel、日志、AI等功能
# 2. package包提供访问当前应用数据的功能,如获取元素、访问全局变量、获取资源文件等功能
# 3. 当此模块作为流程独立运行时执行main函数
# 4. 可视化流程中可以通过"调用模块"的指令使用此模块

import xbot
from xbot import print, sleep
from .import package
from .package import variables as glv
from PIL import Image
from similarities import ClipSimilarity

# 封装函数来计算两张图片之间的相似度
def calculate_image_similarity(image_path1, image_path2,model_path):
    """
    计算两张图片之间的相似度。

    :param image_path1: 第一张图片的路径
    :param image_path2: 第二张图片的路径
    :return: 两张图片的相似度分数
    """
    # 打开图片
    img1 = Image.open(image_path1)
    img2 = Image.open(image_path2)
    similarity = ClipSimilarity(model_name_or_path=model_path,device="cpu")

    # 计算相似度
    score = similarity.similarity([img1], [img2])[0][0]
    return score


def main(args):
  # 图片路径
    image_path1 = r"D:\yys\Desktop\img\4.jpg"
    image_path2 = r"D:\yys\Desktop\img\0_1.png"
    # 模型路径
    model_path=r"D:\Work\python_project\model\clip-vit-base-patch32"
    # 计算相似度
    similarity_score = calculate_image_similarity(image_path1, image_path2, model_path)
    print(f"两张图像之间的相似性得分为: {similarity_score:.2f}")
    pass

2:图搜图

# 使用提醒:
# 1. xbot包提供软件自动化、数据表格、Excel、日志、AI等功能
# 2. package包提供访问当前应用数据的功能,如获取元素、访问全局变量、获取资源文件等功能
# 3. 当此模块作为流程独立运行时执行main函数
# 4. 可视化流程中可以通过"调用模块"的指令使用此模块

import xbot
from xbot import print, sleep
from .import package
from .package import variables as glv
import os
from PIL import Image
from similarities import ClipSimilarity

# 封装函数来实现图搜图的功能
def image_search(query_image_path, corpus_dir_path, topn=1, model_path=None):
    """
    在图库中搜索与查询图片最相似的图片。

    :param query_image_path: 查询图片的路径
    :param corpus_dir_path: 图库的目录路径
    :param topn: 返回最相似的图片数量,默认1 ,返回最相似的图片
    :param model_path: CLIP模型的路径
    :return: 最相似的图片路径列表
    """
    # 打开查询图片
    query_img = Image.open(query_image_path)

    # 获取图库目录下的所有图片文件
    corpus_image_paths = [os.path.join(corpus_dir_path, f) for f in os.listdir(corpus_dir_path)
                          if os.path.isfile(os.path.join(corpus_dir_path, f)) and f.lower().endswith(('.png', '.jpg', '.jpeg', '.gif'))]

    # 加载图库图片
    corpus_imgs = [Image.open(path) for path in corpus_image_paths]

    # 初始化CLIP相似度模型
    similarity = ClipSimilarity(model_name_or_path=model_path,device="cpu") 

    # 添加图库到模型
    similarity.add_corpus(corpus_imgs)

    # 在图库中搜索最相似的图片
    res = similarity.most_similar([query_img], topn=topn)

    # 提取最相似的图片路径
    similar_image_paths = [corpus_image_paths[corpus_id] for corpus_id, _ in res[0].items()]
    return similar_image_paths




def main(args):
    # 使用示例

    # 需要查询图片路径
    query_image_path = r"D:\yys\Desktop\img\0_1.png"

    # 图库的目录路径
    corpus_dir_path = r"D:\yys\Desktop\img"

    # 自定义模型路径
    model_path = r'D:\Work\python_project\model\clip-vit-base-patch32'

    # 进行图搜图
    similar_images = image_search(query_image_path, corpus_dir_path, topn=3, model_path=model_path)

    # 打印搜索结果
    print("相似的图片路径:")
    for path in similar_images:
        print(path)
    pass

3:文搜图


# 使用提醒:
# 1. xbot包提供软件自动化、数据表格、Excel、日志、AI等功能
# 2. package包提供访问当前应用数据的功能,如获取元素、访问全局变量、获取资源文件等功能
# 3. 当此模块作为流程独立运行时执行main函数
# 4. 可视化流程中可以通过"调用模块"的指令使用此模块

import xbot
from xbot import print, sleep
from .import package
from .package import variables as glv

import os
from PIL import Image
from similarities import ClipSimilarity

# 封装函数来实现文搜图的功能
def text_search(query_text, corpus_dir_path, topn=1, model_path=None):
    """
    在图库中搜索与查询文本最相似的图片。

    :param query_text: 查询文本
    :param corpus_dir_path: 包含图库图片的目录路径
    :param topn:  返回最相似的图片数量,默认1 ,返回最相似的图片
    :param model_path: CLIP模型的路径
    :return: 最相似的图片路径列表
    """
    # 获取图库目录下的所有图片文件
    corpus_image_paths = [os.path.join(corpus_dir_path, f) for f in os.listdir(corpus_dir_path)
                          if os.path.isfile(os.path.join(corpus_dir_path, f)) and f.lower().endswith(('.png', '.jpg', '.jpeg', '.gif'))]

    # 加载图库图片
    corpus_imgs = [Image.open(path) for path in corpus_image_paths]

    # 初始化CLIP相似度模型
    similarity = ClipSimilarity(model_name_or_path=model_path,device="cpu") 

    # 添加图库到模型
    similarity.add_corpus(corpus_imgs)

    # 在图库中搜索最相似的图片
    res = similarity.most_similar([query_text], topn=topn)

    # 提取最相似的图片路径
    similar_image_paths = [corpus_image_paths[corpus_id] for corpus_id, _ in res[0].items()]
    return similar_image_paths




def main(args):
    # 使用示例

    # 查询文本
    query_text = '在跑道上的黄色的跑车'

  # 图库图片的目录路径
    corpus_dir_path = r"D:\yys\Desktop\img"

    # 自定义模型路径
    model_path = r'D:\Work\python_project\model\chinese-clip-vit-base-patch16'

    # 进行文搜图
    similar_images = text_search(query_text, corpus_dir_path, topn=1, model_path=model_path)

    # 打印搜索结果
    print("相似的图片路径:")
    for path in similar_images:
        print(path)
    pass

效果预览:







模型下载(模型有两个,根据情况有细微差异):

百度盘:  https://pan.baidu.com/s/1i3FXRlH1tRdU5075qGRVVA?pwd=2dvz /  https://pan.baidu.com/s/1xOqcF1FpYe-YnBEmw8wLFg?pwd=24hw

也可以自己去Huggingface 搜索  OFA-Sys/chinese-clip-vit-huge-patch14 /下载

下载好的模型找个地方解压好,使用的时候需要填模型的路径

收藏8
全部评论1
最新
发布评论
评论