一、安装依赖
pip install ftfy regex tqdm # 安装 clip 的依赖包
pip install git+https://github.com/openai/CLIP.git # 安装 git和clip
二、准备一张图
cat.png
三、完整代码
import torch
# pip install ftfy regex tqdm # 安装 clip 的依赖包
# pip install git+https://github.com/openai/CLIP.git # 安装 git和clip
import clip
from PIL import Image
device = "cuda" if torch.cuda.is_available() else "cpu" # if else 三目运算符
# 加载预训练好的模型
model, preprocess = clip.load("ViT-B/32", device=device) # 选择模型,ViT-B/32是一个基于ViT的模型
# 读取图片和候选类别标签
image = preprocess(Image.open("img/cat.png")).unsqueeze(0).to(
device) # 读取图片,unsqueeze(0)是为了增加一个维度,to(device)是为了将数据传输到GPU或CPU
text = clip.tokenize(["a dog", "a cat", "two cat", "a bike"]).to(device) # 读取文本
with torch.no_grad(): # 以下代码不会计算梯度
# 计算每一张图像和每一个文本的相似度值
logits_per_image, logits_per_text = model(image, text) # 输入图片和文本,输出图片和文本的相似度值
# 对该image与每一个text的相似度值进行softmax
probs = logits_per_image.softmax(
dim=-1).cpu().numpy() # softmax是一个归一化函数,将值映射到0-1之间,且和为1;dim=-1表示对最后一个维度进行softmax; cpu()是为了将数据传输到CPU上,numpy()是为了将数据转换为numpy数组
print("Label probs1:", probs)
四、运行结果
Connected to pydev debugger (build 233.14475.56)
Label probs1: [[1.0730797e-02 9.6741587e-01 2.1420015e-02 4.3338872e-04]]
即 [0.0107, 0.967, 0.0214, 0.000433]
可见识别【a cat】和该图片识别相似度很高。
参考:https://blog.csdn.net/qq_40783025/article/details/129972240
您可以选择一种方式赞助本站
支付宝扫一扫赞助
微信钱包扫描赞助
赏