Tiny Image Captioning

🚀 轻量级图像描述模型

这是一个基于bert - tiny和vit - small的图像描述模型，仅100mb！它在CPU上运行速度极快，能高效完成图像描述任务。

🚀 快速开始

安装依赖

确保你已经安装了transformers库，若未安装，可以使用以下命令进行安装：

pip install transformers requests pillow

运行示例代码

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

model_path = "cnmoro/tiny-image-captioning"

# load the image captioning model and corresponding tokenizer and image processor
model = VisionEncoderDecoderModel.from_pretrained(model_path)
tokenizer = AutoTokenizer.from_pretrained(model_path)
image_processor = AutoImageProcessor.from_pretrained(model_path)

# preprocess an image
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()

# generate caption - suggested settings
generated_ids = model.generate(
    pixel_values,
    temperature=0.7,
    top_p=0.8,
    top_k=50,
    num_beams=3 # you can use 1 for even faster inference with a small drop in quality
)
generated_text = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]

end = time.time()

print(generated_text)
# a group of people walking in the middle of a city.

print(f"Time taken: {end - start} seconds")
# Time taken: 0.11215853691101074 seconds
# on CPU !

✨ 主要特性

• 轻量级：模型仅100mb，占用资源少。
• 高效推理：在CPU上也能实现快速推理。

📦 安装指南

使用pip安装必要的库：

pip install transformers requests pillow

💻 使用示例

基础用法

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

model_path = "cnmoro/tiny-image-captioning"

# load the image captioning model and corresponding tokenizer and image processor
model = VisionEncoderDecoderModel.from_pretrained(model_path)
tokenizer = AutoTokenizer.from_pretrained(model_path)
image_processor = AutoImageProcessor.from_pretrained(model_path)

# preprocess an image
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()

# generate caption - suggested settings
generated_ids = model.generate(
    pixel_values,
    temperature=0.7,
    top_p=0.8,
    top_k=50,
    num_beams=3 # you can use 1 for even faster inference with a small drop in quality
)
generated_text = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]

end = time.time()

print(generated_text)
# a group of people walking in the middle of a city.

print(f"Time taken: {end - start} seconds")
# Time taken: 0.11215853691101074 seconds
# on CPU !

高级用法

如果你希望进一步提高推理速度，可以将num_beams参数设置为1，但可能会导致生成质量略有下降：

# ... 前面的代码保持不变
generated_ids = model.generate(
    pixel_values,
    temperature=0.7,
    top_p=0.8,
    top_k=50,
    num_beams=1 # 更快的推理速度，但质量略有下降
)
generated_text = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
# ... 后续代码保持不变

📚 详细文档

模型信息

📄 许可证

本项目采用Apache 2.0许可证。