Update image_caption.py

image_caption.py

@@ -4,7 +4,7 @@ import os
 from transformers import VisionEncoderDecoderModel, ViTImageProcessor, AutoTokenizer
 import torch
 from PIL import Image
-
+import io
 
 class Caption:
     def __init__(self):
@@ -19,7 +19,7 @@ class Caption:
         )
 
         # device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        self.device = torch.device("cpu")
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
         self.model.to(self.device)
         self.max_length = 16
         self.num_beams = 4
@@ -45,6 +45,32 @@ class Caption:
         preds = [pred.strip() for pred in preds]
         return preds
     
+    def predict_from_memory(self, image_buffers):
+        images = []
+
+        for image_buffer in image_buffers:
+            # Ensure the buffer is positioned at the start
+            if isinstance(image_buffer, io.BytesIO):
+                image_buffer.seek(0)
+            try:
+                i_image = Image.open(image_buffer)
+                if i_image.mode != "RGB":
+                    i_image = i_image.convert("RGB")
+                images.append(i_image)
+            except Exception as e:
+                print(f"Failed to process image buffer: {str(e)}")
+                continue
+
+        return self.process_images(images)
+    
+    def process_images(self, images):
+        pixel_values = self.feature_extractor(images=images, return_tensors="pt").pixel_values
+        pixel_values = pixel_values.to(self.device)
+        output_ids = self.model.generate(pixel_values, **self.gen_kwargs)
+        preds = self.tokenizer.batch_decode(output_ids, skip_special_tokens=True)
+        preds = [pred.strip() for pred in preds]
+        return preds
+
     def get_args(self):
         parser = argparse.ArgumentParser()
         parser.add_argument( "-i",