Update pages/ImageToText.py

pages/ImageToText.py

@@ -0,0 +1,33 @@
+from PIL import Image
+from transformers import AutoProcessor, AutoTokenizer, AutoImageProcessor, AutoModelForCausalLM, BlipForConditionalGeneration, VisionEncoderDecoderModel
+import torch
+
+vitgpt_processor = AutoImageProcessor.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
+vitgpt_model = VisionEncoderDecoderModel.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
+vitgpt_tokenizer = AutoTokenizer.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+vitgpt_model.to(device)
+
+def generate_caption(processor, model, image, tokenizer=None):
+    
+    inputs = processor(images=image, return_tensors="pt").to(device)
+    generated_ids = model.generate(pixel_values=inputs.pixel_values, max_length=50)
+
+    if tokenizer is not None:
+        generated_caption = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
+    else:
+        generated_caption = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
+    return generated_caption
+
+def generate_captions(image):
+    caption_vitgpt = generate_caption(vitgpt_processor, vitgpt_model, image, vitgpt_tokenizer)
+    return caption_vitgpt
+
+uploaded_file = st.file_uploader("Upload your image")
+if uploaded_file is not None:
+    image = Image.open(uploaded_file)
+    generated_caption = generate_caption(vitgpt_processor, vitgpt_model, image, vitgpt_tokenizer)
+    st.write(generated_caption)
+