Update app.py

app.py

@@ -1,3 +1,5 @@
+import cv2
+import numpy as np
 from PIL import Image
 import streamlit as st
 from transformers import GPT2Tokenizer, GPT2LMHeadModel
@@ -10,6 +12,20 @@ model = VisionEncoderDecoderModel.from_pretrained("nlpconnect/vit-gpt2-image-cap
 extractor = ViTFeatureExtractor.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
 tokeniser = AutoTokenizer.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
 
+
+def preprocess_image(image_path):
+    # Load the image using OpenCV
+    image = cv2.imread(image_path)
+
+    # Resize the image while maintaining the aspect ratio
+    resized_image = cv2.resize(image, (224, 224))
+
+    # Add an extra dimension to represent the batch size (assuming a single image)
+    preprocessed_image = np.expand_dims(resized_image, axis=0)
+
+    return preprocessed_image
+
+
 def generate_captions(image):
     generated_caption = tokeniser.decode(model.generate(extractor(image, return_tensors="pt").pixel_values.to("cpu"))[0])
     sentence = generated_caption
@@ -45,15 +61,18 @@ def main():
 
     # check if file has been uploaded
     if uploaded_file is not None:
+        
+        img = preprocess_image(uploaded_file)
+        
         # load the image
-        image = Image.open(uploaded_file).convert("RGB")
+        #image = Image.open(uploaded_file).convert("RGB")
 
         # context as prompt
-        prompt = generate_captions(image)
+        prompt = generate_captions(img)
         st.write("The Context is:", prompt)
 
         # display the image
-        st.image(image)
+        st.image(img)
     
     # Generate button
     if st.button("Generate"):