基于Chinese-Clip的以文搜图和以图搜图实现方案

基于Chinese-Clip的以文搜图和以图搜图实现方案

引用

CSDN

1.

https://blog.csdn.net/qq_44908396/article/details/144537426

本文介绍了一种基于Chinese-Clip的以文搜图和以图搜图实现方案。通过详细的技术讲解和代码示例，帮助读者快速掌握这一前沿技术的应用方法。

前言

目前网上能够找到的资料有限，要么收费，要么配置复杂。本文作者决定自己动手实现一个demo，功能清单受启发于Nvidia AI lab实验室的nanodb项目，旨在开发一个可以实现以文搜图和以图搜图的demo。由于作者非科班出身，代码知识面较窄，因此未实现网页功能，仅提供demo，具体业务功能需要自行编写。

实现思路

整体思路如下图所示，我们先将图像使用clip生成其对应的特征向量存入数据库当中，然后通过图像输入或者文本输入进行查询，需要注意，图像和文本输入有一项即可。

环境配置

• Chinese-Clip: GitHub - OFA-Sys/Chinese-CLIP
• Milvus:

  pip install -U pymilvus
  

• Pytorch:

  pip install torch==1.13.0+cu117 torchvision==0.14.0+cu117 torchaudio==0.13.0 --extra-index-url https://download.pytorch.org/whl/cu117
  

源码

数据入库

import torch
from PIL import Image
import cn_clip.clip as clip
from cn_clip.clip import load_from_name, available_models
import numpy as np
import os
from pymilvus import MilvusClient

client = MilvusClient("/home/project_python/Chinese-CLIP/my_database/coco2017.db")
if client.has_collection(collection_name="text_image"):
    client.drop_collection(collection_name="text_image")
client.create_collection(
    collection_name="text_image",
    dimension=512,  # The vectors we will use in this demo has 768 dimensions
    metric_type="COSINE"
)
print("Available models:", available_models())

def getFileList(dir, Filelist, ext=None):
    """
    获取文件夹及其子文件夹中文件列表
    输入 dir：文件夹根目录
    输入 ext: 扩展名
    返回： 文件路径列表
    """
    newDir = dir
    if os.path.isfile(dir):
        if ext is None:
            Filelist.append(dir)
        else:
            if ext in dir:
                Filelist.append(dir)
    elif os.path.isdir(dir):
        for s in os.listdir(dir):
            newDir = os.path.join(dir, s)
            getFileList(newDir, Filelist, ext)
    return Filelist

if __name__ == "__main__":
    device = "cuda" if torch.cuda.is_available() else "cpu"
    model, preprocess = load_from_name("ViT-B-16", device=device, download_root='./')
    model.eval()
    img_dir = r"/home/project_python/Chinese-CLIP/my_dataset/coco"
    image_path_list = []
    image_path_list = getFileList(img_dir, image_path_list, '.jpg')
    data = []
    i = 0
    for image_path in image_path_list:
        temp = {}
        image = preprocess(Image.open(image_path)).unsqueeze(0).to(device)
        with torch.no_grad():
            image_features = model.encode_image(image)
            # 对特征进行归一化，请使用归一化后的图文特征用于下游任务
            image_features /= image_features.norm(dim=-1, keepdim=True)
            image_features = image_features.cpu().numpy().astype(np.float32).flatten()
        # 将特征向量转换为字符串
        #features_str = ','.join(map(str, image_features.flatten()))
        temp['id'] = i
        temp['image_path'] = image_path
        temp['vector'] = image_features
        data.append(temp)
        i = i + 1
        print(i)
    res = client.insert(collection_name="text_image", data=data)

上述代码会在指定路径生成一个coco2017.db的文件，这就说明数据完成了入库，我们接下来进行调用。

数据查询

import torch
from PIL import Image,ImageDraw, ImageFont
import cn_clip.clip as clip
from cn_clip.clip import load_from_name, available_models
import numpy as np
import time
from pymilvus import MilvusClient

client = MilvusClient("/home/project_python/Chinese-CLIP/my_database/coco2017.db")
print("Available models:", available_models())
# Available models: ['ViT-B-16', 'ViT-L-14', 'ViT-L-14-336', 'ViT-H-14', 'RN50']

def display_single_image_with_text(image_path):
    with Image.open(image_path) as img:
        draw = ImageDraw.Draw(img)
        # 设置字体和字号，这里假设你有一个可用的字体文件，例如 Arial.ttf
        # 如果没有，可以使用系统默认字体
        try:
            font = ImageFont.truetype("Arial.ttf", 30)
        except IOError:
            font = ImageFont.load_default()
        # 文本内容和颜色
        text = "Example image"
        text_color = (255, 0, 0)  # 红色
        # 文本位置
        text_position = (10, 10)
        # 绘制文本
        draw.text(text_position, text, fill=text_color, font=font)
        # 显示图像
        img.show()

def display_images_in_grid(image_paths, images_per_row=3):
    # 计算需要的行数
    num_images = len(image_paths)
    num_rows = (num_images + images_per_row - 1) // images_per_row
    # 打开所有图像并调整大小
    images = []
    for path in image_paths:
        with Image.open(path) as img:
            img = img.resize((200, 200))  # 调整图像大小以适应画布
            images.append(img)
    # 创建一个空白画布
    canvas_width = images_per_row * 200
    canvas_height = num_rows * 200
    canvas = Image.new('RGB', (canvas_width, canvas_height), (255, 255, 255))
    # 将图像粘贴到画布上
    for idx, img in enumerate(images):
        row = idx // images_per_row
        col = idx % images_per_row
        position = (col * 200, row * 200)
        canvas.paste(img, position)
    # 显示画布
    canvas.show()

def load_model(device):
    model, preprocess = load_from_name("ViT-B-16", device=device, download_root='./')
    model.eval()
    return model, preprocess

def text_encode(text,device):
    new_text = clip.tokenize([text]).to(device)
    with torch.no_grad():
        text_features = model.encode_text(new_text)
        text_features /= text_features.norm(dim=-1, keepdim=True)
        text_features = text_features.cpu().numpy().astype(np.float32)
    return text_features

def image_encode(model,preprocess,image_path,device):
    image = preprocess(Image.open(image_path)).unsqueeze(0).to(device)
    with torch.no_grad():
        image_features = model.encode_image(image)
        image_features /= image_features.norm(dim=-1, keepdim=True)
        image_features = image_features.cpu().numpy().astype(np.float32)
    return image_features

if __name__ == "__main__":
    search_text = "猫"
    search_image_path = "/home/project_python/Chinese-CLIP/my_dataset/coco/val2017/000000000285.jpg"
    device = "cuda" if torch.cuda.is_available() else "cpu"
    model, preprocess = load_model(device)
    text_flag = False
    if text_flag:
        text_features = text_encode(search_text,device)
        results = client.search(
            "text_image",
            data=text_features,
            output_fields=["image_path"],
            search_params={"metric_type": "COSINE"},
            limit=36
        )
    else:
        display_single_image_with_text(search_image_path)
        image_features = image_encode(model,preprocess,search_image_path,device)
        results = client.search(
            "text_image",
            data=image_features,
            output_fields=["image_path"],
            search_params={"metric_type": "COSINE"},
            limit=36
        )
    image_list = []
    for i,result in enumerate(results[0]):
        image_list.append(result["entity"]["image_path"])
    display_images_in_grid(image_list,9)

上述代码使用text_flag控制是以文搜图还是以图搜图，True时为以文搜图，False时为以图搜图。

实现效果

以文搜图

以图搜图

附加

1. Chinese-Clip开放了onnx和trt推理，可以根据自己的需求改进，参考链接如下：

• FoundationModel/Chinese-CLIP

2. 可以根据自己的时间安排复现下面的项目，作者这里只是demo，上不了台面，参考链接如下：

• nanoDB