Update image_to_text.py

image_to_text.py

@@ -2,6 +2,11 @@ from PIL import Image
 import torch
 from transformers import VisionEncoderDecoderModel, ViTImageProcessor, AutoTokenizer
 
+def preprocess_image(image, target_size=(224, 224)):
+    # Resize the image to a fixed size
+    image = image.resize(target_size, Image.LANCZOS)
+    return image
+
 def generate_caption(image):
     # Load pre-trained model and tokenizer
     model = VisionEncoderDecoderModel.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
@@ -12,13 +17,14 @@ def generate_caption(image):
     model.to(device)
 
     # Prepare image
-    if isinstance(image, Image.Image):
-        image = image.convert('RGB')
-    else:
-        image = Image.open(image).convert('RGB')
+    if not isinstance(image, Image.Image):
+        image = Image.open(image)
+    
+    image = image.convert('RGB')
+    image = preprocess_image(image)
     
-    # Add padding=True to handle images of different sizes
-    pixel_values = feature_extractor(images=[image], return_tensors="pt", padding=True).pixel_values
+    # Convert image to pixel values
+    pixel_values = feature_extractor(images=image, return_tensors="pt").pixel_values
     pixel_values = pixel_values.to(device)
 
     # Generate caption