库名称：transformers
许可证：mit
数据集：

• armaggheddon97/lego_brick_captions
  语言：
• en
  基础模型：
• openai/clip-vit-base-patch32
  管道标签：zero-shot-classification

clip-vit-base-patch32_lego-brick 模型卡

模型详情

此模型是基于 lego_brick_captions 数据集对 openai/clip-vit-base-patch32 CLIP（对比语言-图像预训练）模型进行微调的版本，专门用于匹配乐高积木图像与其对应的文本描述。

[!NOTE]
若对代码感兴趣，请参考我的 GitHub 上的微调脚本。

🔍 探索此模型的强大功能

是否曾为找不到某块乐高积木的名称而烦恼？或者虽有模糊印象或图片，却无法确定具体零件编号？BricksFinder 正是为此而生！

输入描述如“蓝色弧形斜坡”或上传积木图片，模型将施展魔法，为您找到最接近的匹配项。它会展示一系列与您设想（甚至更优）的积木图片！

无论是乐高爱好者、建造者，还是喜欢在积木中寻宝的人，这都是完美工具。快来 Colab 上体验实时演示吧！

模型描述

• 开发者： 基础模型由 OpenAI 开发，微调模型由 armaggheddon97 完成。
• 模型类型： CLIP（对比语言-图像预训练）模型。
• 语言： 模型输入需为英文文本。
• 许可证： MIT 许可证。
• 微调自 clip-vit-base-patch32： 此模型是在 lego_brick_captions 数据集上对 openai/clip-vit-base-patch32 进行微调的版本。模型在数据集的 80-20 训练-验证分割上微调了 7 个周期。微调脚本详情请见 GitHub。

使用 🤗 transformers

• 加载模型与处理器：

  import torch  
  from transformers import CLIPProcessor, CLIPModel  
  
  device = "cuda" if torch.cuda.is_available() else "cpu"  
  
  model = CLIPModel.from_pretrained("armaggheddon97/clip-vit-base-patch32_lego-brick", device_map="auto").to(device)  
  processor = CLIPProcessor.from_pretrained("armaggheddon97/clip-vit-base-patch32_lego-brick", device_map="auto").to(device)  
  

• 使用 Auto 类：

  from transformers import AutoModelForZeroShotImageClassification, AutoProcessor  
  
  model = AutoModelForZeroShotImageClassification.from_pretrained("armaggheddon97/clip-vit-base-patch32_lego-brick")  
  processor = AutoProcessor.from_pretrained("armaggheddon97/clip-vit-base-patch32_lego-brick")  
  

• 使用 pipeline：

  from transformers import pipeline  
  
  model = "armaggheddon97/clip-vit-base-patch32_lego-brick"  
  clip_classifier = pipeline("zero-shot-image-classification", model=model)  
  

以 float16 精度加载

默认模型为 float32 精度。如需以 float16 精度加载以加速推理，可使用以下代码：

from transformers import CLIPProcessor, CLIPModel  

model = CLIPModel.from_pretrained("armaggheddon97/clip-vit-base-patch32_lego-brick", dtype=torch.float16)  
processor = CLIPProcessor.from_pretrained("armaggheddon97/clip-vit-base-patch32_lego-brick")  

或直接使用 torch：

import torch  
from transformers import CLIPModel  

model = CLIPModel.from_pretrained("armaggheddon97/clip-vit-base-patch32_lego-brick")  
model_fp16 = model.to(torch.float16)  

应用场景

生成嵌入

• 仅嵌入文本：

  import torch  
  from transformers import CLIPTokenizerFast, CLIPModel  
  
  device = "cuda" if torch.cuda.is_available() else "cpu"  
  
  model = CLIPModel.from_pretrained("armaggheddon97/clip-vit-base-patch32_lego-brick", device_map="auto").to(device)  
  tokenizer = CLIPTokenizerFast.from_pretrained("armaggheddon97/clip-vit-base-patch32_lego-brick")  
  
  text = ["a photo of a lego brick"]  
  tokens = tokenizer(text, return_tensors="pt", padding=True).to(device)  
  outputs = model.get_text_features(**tokens)  
  

• 仅嵌入图像：

  import torch  
  from PIL import Image  
  from transformers import CLIPProcessor, CLIPModel  
  
  device = "cuda" if torch.cuda.is_available() else "cpu"  
  
  model = CLIPModel.from_pretrained("armaggheddon97/clip-vit-base-patch32_lego-brick", device_map="auto").to(device)  
  processor = CLIPProcessor.from_pretrained("armaggheddon97/clip-vit-base-patch32_lego-brick", device_map="auto").to(device)  
  
  image = Image.open("path_to_image.jpg")  
  inputs = processor(images=image, return_tensors="pt").to(device)  
  outputs = model.get_image_features(**inputs)  
  

零样本图像分类

import torch  
from PIL import Image  
from transformers import CLIPProcessor, CLIPModel  
from datasets import load_dataset  

device = "cuda" if torch.cuda.is_available() else "cpu"  

model = CLIPModel.from_pretrained("armaggheddon97/clip-vit-base-patch32_lego-brick", device_map="auto").to(device)  
processor = CLIPProcessor.from_pretrained("armaggheddon97/clip-vit-base-patch32_lego-brick", device_map="auto").to(device)  

dataset = load_dataset("armaggheddon97/lego_brick_captions", split="test")  

captions = [  
    "a photo of a lego brick with a 2x2 plate",  
    "a photo of gray minifigure legs",  
    "a photo of a brick with a curved slope",  
]  
image = dataset[0]["image"]  

inputs = processor(text=captions, images=image, return_tensors="pt", padding=True).to(device)  
outputs = model(**inputs)  

logits_per_image = outputs.logits_per_image  
probabilities = logits_per_image.softmax(dim=1)  
max_prob_idx = torch.argmax(logits_per_image, dim=1)  

结果

目标是通过微调使模型能更准确地区分基于文本描述的积木图像。在准确率方面，微调模型与基础模型表现相近。但在分类任务测试中（如零样本图像分类部分代码所示），微调模型能以更高置信度准确分类图像。

例如，测试输入为：

• 沙绿色 2x2 人仔腿部零件，顶部有两个轴孔。腿部印有浅灰、橙色和米色交织的布料图案。零件为实心，除轴孔外无其他连接点。
• 中绿色 1x1 圆形人仔头部，印有黄色眼睛、粉色花朵元素和露齿笑容。顶部为标准凸点连接，无其他可见连接点。印刷细节覆盖大部分表面。
• 白色 2x2 砖块，四个凸点均印有乐高标志。砖块为标准尺寸，无额外孔洞或特征。颜色为明亮的偏白色。

输入图像为：

微调模型输出：

• 100.00%: "沙绿色 2x2 人仔腿部零件..."
• 0.00%: "中绿色 1x1 圆形人仔头部..."
• 0.00%: "白色 2x2 砖块..."

而基础模型输出：

• 98.7%: "沙绿色 2x2 人仔腿部零件..."
• 1.24%: "中绿色 1x1 圆形人仔头部..."
• 0.00%: "白色 2x2 砖块..."

这表明微调模型能更精准地基于文本描述分类图像，而基础模型虽正确率相近，但置信度略低。

在整个数据集上运行相同任务（1 条正确描述 + 2 条随机采样描述），结果如下：

热图展示了微调与基础模型的归一化文本逻辑值：

• 输入： 每个样本包含一张乐高积木图像及三条描述：
  • 正确描述（位置 0）。
  • 两条随机错误描述（位置 1 和 2）。
• 输出： 模型为每条描述生成文本逻辑值，反映图像嵌入与文本嵌入的相似度。逻辑值经归一化以便可视化。
• 热图说明：
  • 每行代表一个样本。
  • 每列代表一条描述（正确描述在第 0 列）。
  • 颜色深浅表示模型对描述的置信度，对比度越大效果越好。

基础模型（右图）对所有类别均未显示高置信度，区分能力较弱（逻辑值差异小），但准确率仍达 97.46%。
微调模型（左图）对正确描述的置信度显著更高，准确率达 99.23%。

基于 short_caption 的微调

作为实验，模型还在数据集的 short_caption 列上进行了微调。与基础模型及基于 caption 列微调的模型对比，结果如下：

基于 short_caption 微调的模型：

• 100.00%: "深棕腿部与长袍及深橙腰带印刷"
• 0.00%: "黄色眼睛、尖牙气泡印刷的人仔头部"
• 0.00%: "无内脊 2x2 砖块"

基于 caption 微调的模型：

• 100.00%: "深棕腿部与长袍及深橙腰带印刷"
• 0.00%: "黄色眼睛、尖牙气泡印刷的人仔头部"
• 0.00%: "无内脊 2x2 砖块"

基础模型：

• 0.00%: "深棕腿部与长袍及深橙腰带印刷"
• 22.07%: "黄色眼睛、尖牙气泡印刷的人仔头部"
• 77.79%: "无内脊 2x2 砖块"

尽管基于 short_caption 微调，其与基于 caption 微调的模型结果高度相似，仅置信度跨度更大。基础模型在此任务中表现显著较差。

在整个数据集上测试时：

基于 short_caption 微调的模型准确率达 99.99%，而基于 caption 的模型为 98.48%。虽前者更准确，但后者因 caption 列提供的文本灵活性更优，故最终选择上传后者。

基础模型在全数据集测试中准确率仍约 97%，说明此前测试样本可能为其异常案例。