Create app.py

app.py

@@ -0,0 +1,40 @@
+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])
+
+