基于电商图文数据跨模态检索 图文互检系统

跨模态检索 图文互检系统

可替换泰迪杯数据和模型权重（或者其他场景数据集）进行推理互检，实现图文互检、图图互检、文图互检三个功能以及模型特征向量输出，可做毕设。

效果预览：

视频预览：
bilibili视频演示

本地启动：

1、启动flask服务

2、启动app

实现原理：

训练后的最优模型pt转为onnx模型：将Pytorch模型checkpoint转换为ONNX格式的代码

cd Chinese-CLIP/
export CUDA_VISIBLE_DEVICES=0
export PYTHONPATH=${PYTHONPATH}:`pwd`/cn_clip

# ${DATAPATH}的指定，请参考Readme"代码组织"部分创建好目录，尽量使用相对路径：https://github.com/OFA-Sys/Chinese-CLIP#代码组织
checkpoint_path=${DATAPATH}/pretrained_weights/clip_cn_vit-b-16.pt # 指定要转换的ckpt完整路径
mkdir -p ${DATAPATH}/deploy/ # 创建ONNX模型的输出文件夹

python cn_clip/deploy/pytorch_to_onnx.py \
       --model-arch ViT-B-16 \
       --pytorch-ckpt-path ${checkpoint_path} \
       --save-onnx-path ${DATAPATH}/deploy/vit-b-16 \
       --convert-text --convert-vision

其中各配置项定义如下：

• model-arch: 模型规模，选项包括["RN50", "ViT-B-16", "ViT-L-14", "ViT-L-14-336", "ViT-H-14"]，
• pytorch-ckpt-path: 指定Pytorch模型ckpt路径，上面的代码示例中我们指定为预训练的ckpt路径，也可以指定为用户finetune ckpt的位置。ckpt中的参数需要与model-arch指定的模型规模对应
• save-onnx-path: 指定输出ONNX格式模型的路径（前缀）。完成转换后，代码将分别输出文本侧和图像侧的ONNX格式编码模型文件，FP32与FP16各一版，该参数即指定了以上输出文件的路径前缀
• convert-text和convert-vision: 指定是否转换文本侧和图像侧模型
• context-length（可选）: 指定文本侧ONNX模型，接收输入的序列长度，默认为我们预训练ckpt所使用的52
• download-root（可选）: 如果不指定pytorch-ckpt-path，代码将根据model-arch自动下载Chinese-CLIP官方预训练ckpt用于转换，存放于download-root指定的目录

运行此代码转换完成后，将得到以下的log输出：

Finished PyTorch to ONNX conversion...
>>> The text FP32 ONNX model is saved at ${DATAPATH}/deploy/vit-b-16.txt.fp32.onnx
>>> The text FP16 ONNX model is saved at ${DATAPATH}/deploy/vit-b-16.txt.fp16.onnx with extra file ${DATAPATH}/deploy/vit-b-16.txt.fp16.onnx.extra_file
>>> The vision FP32 ONNX model is saved at ${DATAPATH}/deploy/vit-b-16.img.fp32.onnx
>>> The vision FP16 ONNX model is saved at ${DATAPATH}/deploy/vit-b-16.img.fp16.onnx with extra file ${DATAPATH}/deploy/vit-b-16.img.fp16.onnx.extra_file

定义模型对象代码。该对象具有图像特征导出、文本特征导出：

import onnxruntime
from PIL import Image
import numpy as np
import torch
from cn_clip.clip.utils import _MODEL_INFO, image_transform
import cn_clip.clip as clip



class OnnxModel():
    def __init__(self, img_onnx_model_path, txt_onnx_model_path, model_arch = "ViT-B-16", ):
        self.img_model_path = img_onnx_model_path
        self.model_arch = model_arch


        # 载入ONNX图像侧模型（**请替换${DATAPATH}为实际的路径**）
        img_sess_options = onnxruntime.SessionOptions()
        img_run_options = onnxruntime.RunOptions()
        img_run_options.log_severity_level = 2
        self.img_session = onnxruntime.InferenceSession(self.img_model_path,
                                                   sess_options=img_sess_options,
                                                   providers=["CUDAExecutionProvider"])

        txt_sess_options = onnxruntime.SessionOptions()
        txt_run_options = onnxruntime.RunOptions()
        txt_run_options.log_severity_level = 2
        self.txt_session = onnxruntime.InferenceSession(txt_onnx_model_path,
                                                   sess_options=txt_sess_options,
                                                   providers=["CUDAExecutionProvider"])

    def extract_image_feat(self,image):
        # 预处理图片
        # model_arch = "ViT-B-16"  # 这里我们使用的是ViT-B-16规模，其他规模请对应修改
        preprocess = image_transform(_MODEL_INFO[self.model_arch]['input_resolution'])
        image = preprocess(image).unsqueeze(0)

        # 用ONNX模型计算图像侧特征
        image_features = self.img_session.run(["unnorm_image_features"], {"image": image.cpu().numpy()})[0]  # 未归一化的图像特征
        image_features = torch.tensor(image_features)
        image_features /= image_features.norm(dim=-1, keepdim=True)  # 归一化后的Chinese-CLIP图像特征，用于下游任务
        print(image_features.shape)  # Torch Tensor shape: [1, 特征向量维度]
        image_features_tensor = image_features.cuda()
        return image_features_tensor

    def extract_text_feat(self,text_):
        # 为4条输入文本进行分词。序列长度指定为52，需要和转换ONNX模型时保持一致（参见转换时的context-length参数）
        text = clip.tokenize(["{}".format(text_)], context_length=52)
        print(text)

        text_features = []
        for i in range(len(text)):
            one_text = np.expand_dims(text[i].cpu().numpy(), axis=0)
            text_feature = self.txt_session.run(["unnorm_text_features"], {"text": one_text})[0]  # 未归一化的文本特征
            text_feature = torch.tensor(text_feature)
            text_features.append(text_feature)
        # 拿到text_feature
        text_features = torch.squeeze(torch.stack(text_features), dim=1)  # 4个特征向量stack到一起
        text_features = text_features / text_features.norm(dim=1, keepdim=True)
        # 转为torch.Tensor 并放在cuda
        text_feat_tensor = text_features.cuda()
        return text_feat_tensor

编写后端模型推理API，初始化模型对象，调用模型进行互检。

项目框架介绍：

代码框架文件夹：

框架结构：

数据集、模型替换：

获取链接

烦请移步社区