| import streamlit as st | |
| from PIL import Image | |
| from transformers import VisionEncoderDecoderModel, ViTImageProcessor, GPT2TokenizerFast | |
| import torch | |
| from PIL import Image | |
| model = VisionEncoderDecoderModel.from_pretrained("nlpconnect/vit-gpt2-image-captioning") | |
| feature_extractor = ViTImageProcessor.from_pretrained("nlpconnect/vit-gpt2-image-captioning") | |
| tokenizer=GPT2TokenizerFast.from_pretrained("nlpconnect/vit-gpt2-image-captioning") | |
| gen_kwargs1 ={"max_length": 4,"num_beams": 2} | |
| gen_kwargs2 ={"max_length": 32,"num_beams": 16} | |
| def predict_step(images): | |
| pixel_values = feature_extractor(images=images, return_tensors='pt').pixel_values | |
| output_ids1 = model.generate(pixel_values) | |
| output_ids2 = model.generate(pixel_values,**gen_kwargs1) | |
| output_ids3 = model.generate(pixel_values,**gen_kwargs2) | |
| preds1 = tokenizer.batch_decode(output_ids1, skip_special_tokens=True) | |
| preds2 = tokenizer.batch_decode(output_ids2, skip_special_tokens=True) | |
| preds3 = tokenizer.batch_decode(output_ids3, skip_special_tokens=True) | |
| preds1 =[pred.strip() for pred in preds1] | |
| preds2 =[pred.strip() for pred in preds2] | |
| preds3 =[pred.strip() for pred in preds3] | |
| return preds1[0],preds2[0],preds3[0] | |
| st.title("Image Caption Generator") | |
| upload_image = st.file_uploader(label='Upload image', type=['png', 'jpg','jpeg'], accept_multiple_files=False) | |
| if upload_image is not None: | |
| image = Image.open(upload_image) | |
| if image.mode != "RGB": | |
| image = image.convert(mode="RGB") | |
| output = predict_step([image]) | |
| st.header("Captions are : ") | |
| st.text(output[0]) | |
| st.text(output[1]) | |
| st.text(output[2]) | |