本文不生產技術，只做技術的搬運工！

前言

????????作者昨天使用cn_clip庫實現了一版，但是覺得大家復現配置環境可能有點復雜，因此有使用transformers庫實現了一版，提供大家選擇，第一篇參考鏈接如下：

Chinese-Clip實現以文搜圖和以圖搜圖-CSDN博客文章瀏覽閱讀728次，點贊9次，收藏17次。使用clip實現以文搜圖和以圖搜圖的圖文檢索功能https://blog.csdn.net/qq_44908396/article/details/144537426

?環境配置

transformers:

pip install transformers

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

源碼

數據入庫

from PIL import Image
import requests
from transformers import ChineseCLIPProcessor, ChineseCLIPModel
import torch
import os
import numpy as np
from pymilvus import MilvusClient
client = MilvusClient("BlingPic.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 dimensionsmetric_type="COSINE"
)def getFileList(dir, Filelist, ext=None):"""獲取文件夾及其子文件夾中文件列表輸入 dir：文件夾根目錄輸入 ext: 擴展名返回： 文件路徑列表"""newDir = dirif 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 Filelistif __name__ == "__main__":device = "cuda" if torch.cuda.is_available() else "cpu"model = ChineseCLIPModel.from_pretrained("OFA-Sys/chinese-clip-vit-base-patch16")model.to(device)preprocess = ChineseCLIPProcessor.from_pretrained("OFA-Sys/chinese-clip-vit-base-patch16")model.eval()img_dir = r"/home/turing/圖片/BlingPic"image_path_list = []image_path_list = getFileList(img_dir, image_path_list, '.jpg')data = []i = 0for image_path in image_path_list:temp = {}image = Image.open(image_path)with torch.no_grad():inputs = preprocess(images=image, return_tensors="pt").to(device)image_features = model.get_image_features(**inputs)image_features = image_features / image_features.norm(dim=-1, keepdim=True)  # normalizeimage_features = image_features.cpu().numpy().astype(np.float32).flatten()# 將特征向量轉換為字符串temp['id'] = itemp['image_path'] = image_pathtemp['vector'] = image_featuresdata.append(temp)i = i + 1print(i)res = client.insert(collection_name="text_image", data=data)

上述代碼會在指定路徑生成一個BlingPic.db的文件，這就說明數據完成了入庫，我們接下來進行調用

數據查詢

from PIL import Image,ImageDraw,ImageFont
from transformers import ChineseCLIPProcessor, ChineseCLIPModel,AutoTokenizer
import torch
import numpy as np
from pymilvus import MilvusClient
client = MilvusClient("BlingPic.db")
# 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 * 200canvas_height = num_rows * 200canvas = Image.new('RGB', (canvas_width, canvas_height), (255, 255, 255))# 將圖像粘貼到畫布上for idx, img in enumerate(images):row = idx // images_per_rowcol = idx % images_per_rowposition = (col * 200, row * 200)canvas.paste(img, position)# 顯示畫布canvas.show()def load_model(device):model = ChineseCLIPModel.from_pretrained("OFA-Sys/chinese-clip-vit-base-patch16")model.to(device)preprocess = ChineseCLIPProcessor.from_pretrained("OFA-Sys/chinese-clip-vit-base-patch16")model.eval()return model, preprocessdef text_encode(model,text,device):tokenizer = AutoTokenizer.from_pretrained("OFA-Sys/chinese-clip-vit-base-patch16")inputs = tokenizer(text, return_tensors="pt").to(device)with torch.no_grad():text_features = model.get_text_features(**inputs)text_features /= text_features.norm(dim=-1, keepdim=True)text_features = text_features.cpu().numpy().astype(np.float32)return text_featuresdef image_encode(model,preprocess,image_path,device):image = Image.open(image_path)with torch.no_grad():inputs = preprocess(images=image, return_tensors="pt").to(device)image_features = model.get_image_features(**inputs)image_features = image_features / image_features.norm(dim=-1, keepdim=True)  # normalizeimage_features = image_features.cpu().numpy().astype(np.float32)return image_featuresif __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 = Falseif text_flag:text_features = text_encode(model,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時為以圖搜圖

實現效果

以文搜圖

以圖搜圖

示例圖像：

搜索結果：

附加

權重下載遇到問題參考如下鏈接：

解決OSError: We couldn‘t connect to ‘https://huggingface.co‘ to load this file-CSDN博客文章瀏覽閱讀1.4k次，點贊6次，收藏2次。解決hugging face無法下載模型的問題https://blog.csdn.net/qq_44908396/article/details/142516867