Update image_to_text.py

image_to_text.py

@@ -1,11 +1,16 @@
 from PIL import Image
 import torch
+import numpy as np
 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
+    # Convert to numpy array and normalize
+    image_array = np.array(image) / 255.0
+    # Transpose to (channels, height, width) format
+    image_array = np.transpose(image_array, (2, 0, 1))
+    return image_array
 
 def generate_caption(image):
     # Load pre-trained model and tokenizer
@@ -21,10 +26,13 @@ def generate_caption(image):
         image = Image.open(image)
     
     image = image.convert('RGB')
-    image = preprocess_image(image)
+    image_array = preprocess_image(image)
     
-    # Convert image to pixel values
-    pixel_values = feature_extractor(images=image, return_tensors="pt").pixel_values
+    # Create a batch with a single image
+    batch = np.expand_dims(image_array, axis=0)
+    
+    # Convert to tensor
+    pixel_values = torch.tensor(batch).float()
     pixel_values = pixel_values.to(device)
 
     # Generate caption