| from transformers import VisionEncoderDecoderModel, ViTImageProcessor, AutoTokenizer |
| from PIL import Image |
| import torch |
|
|
| model = VisionEncoderDecoderModel.from_pretrained("nlpconnect/vit-gpt2-image-captioning") |
| processor = ViTImageProcessor.from_pretrained("nlpconnect/vit-gpt2-image-captioning") |
| tokenizer = AutoTokenizer.from_pretrained("nlpconnect/vit-gpt2-image-captioning") |
|
|
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
| model.to(device) |
|
|
|
|
| def generate_caption(image: Image.Image) -> str: |
| if image.mode != "RGB": |
| image = image.convert(mode="RGB") |
|
|
| pixel_values = processor(images=image, return_tensors="pt").pixel_values.to(device) |
| output_ids = model.generate(pixel_values, max_length=16, num_beams=1) |
| caption = tokenizer.decode(output_ids[0], skip_special_tokens=True) |
| return caption.strip() |
|
|
