Nano Image Captioning

🚀 納米圖像字幕生成模型

這是一個基於BERT-Tiny和ViT-Tiny的圖像字幕生成模型，僅40MB！它在CPU上也能快速運行，為圖像添加描述信息提供了高效解決方案。

🚀 快速開始

此圖像字幕生成模型能快速為圖像生成描述。以下是使用該模型的步驟：

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

model_path = "cnmoro/nano-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)
# a group of people are in the middle of a city.

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

💻 使用示例

基礎用法

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

model_path = "cnmoro/nano-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)
# a group of people are in the middle of a city.

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

高級用法

# 如果你需要在多個圖像上進行批量處理，可以將圖像URL存儲在列表中，然後循環處理每個圖像。
from transformers import AutoTokenizer, AutoImageProcessor, VisionEncoderDecoderModel
import requests, time
from PIL import Image

model_path = "cnmoro/nano-image-captioning"

# 加載圖像字幕生成模型以及對應的分詞器和圖像處理器
model = VisionEncoderDecoderModel.from_pretrained(model_path)
tokenizer = AutoTokenizer.from_pretrained(model_path)
image_processor = AutoImageProcessor.from_pretrained(model_path)

image_urls = [
    "https://upload.wikimedia.org/wikipedia/commons/thumb/4/47/New_york_times_square-terabass.jpg/800px-New_york_times_square-terabass.jpg",
    "https://example.com/another_image.jpg"
]

for url in image_urls:
    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"Time taken: {end - start} seconds")

📄 許可證

本項目採用Apache-2.0許可證。

📚 詳細文檔