许可证：apache-2.0
语言：

• 英文
  基础模型：
• google/bert_uncased_L-4_H-256_A-4
• WinKawaks/vit-small-patch16-224
  任务标签：图像转文本
  库名称：transformers
  标签：
• vit
• bert
• 视觉
• 字幕
• 字幕生成
• 图像

一个基于bert-mini和vit-small的图像字幕生成模型，仅重130MB！

在CPU上运行速度极快。

from transformers import AutoTokenizer, AutoImageProcessor, VisionEncoderDecoderModel
import requests, time
from PIL import Image

model_path = "cnmoro/mini-image-captioning"

# 加载图像字幕生成模型及对应的分词器和图像处理器
model = VisionEncoderDecoderModel.from_pretrained(model_path)
tokenizer = AutoTokenizer.from_pretrained(model_path)
image_processor = AutoImageProcessor.from_pretrained(model_path)

# 预处理图像
url = "https://upload.wikimedia.org/wikipedia/commons/thumb/4/47/New_york_times_square-terabass.jpg/800px-New_york_times_square-terabass.jpg"
image = Image.open(requests.get(url, stream=True).raw)
pixel_values = image_processor(image, return_tensors="pt").pixel_values

start = time.time()

# 生成字幕 - 推荐参数设置
generated_ids = model.generate(
    pixel_values,
    temperature=0.7,
    top_p=0.8,
    top_k=50,
    num_beams=3 # 可使用1以更快推理，但质量略有下降
)
generated_text = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]

end = time.time()

print(generated_text)
# 一大群人穿过繁华的城市。

print(f"耗时: {end - start} 秒")
# 耗时: 0.19002342224121094 秒
# 在CPU上！