Update image_caption.py

image_caption.py

@@ -4,26 +4,31 @@ import os
 from transformers import VisionEncoderDecoderModel, ViTImageProcessor, AutoTokenizer
 import torch
 from PIL import Image
-
+import io
+import google.generativeai as genai
 
 class Caption:
     def __init__(self):
-        self.model = VisionEncoderDecoderModel.from_pretrained(
-            "nlpconnect/vit-gpt2-image-captioning"
-        )
-        self.feature_extractor = ViTImageProcessor.from_pretrained(
-            "nlpconnect/vit-gpt2-image-captioning"
-        )
-        self.tokenizer = AutoTokenizer.from_pretrained(
-            "nlpconnect/vit-gpt2-image-captioning"
-        )
 
-        # device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        self.device = torch.device("cpu")
-        self.model.to(self.device)
-        self.max_length = 16
-        self.num_beams = 4
-        self.gen_kwargs = {"max_length": self.max_length, "num_beams": self.num_beams}
+        self.api_key = 'AIzaSyAFG94rVbm9eWepO5uPGsMha8XJ-sHbMdA'
+        genai.configure(api_key=self.api_key)
+        self.model = genai.GenerativeModel(model_name="gemini-pro-vision")
+    #     self.model = VisionEncoderDecoderModel.from_pretrained(
+    #         "nlpconnect/vit-gpt2-image-captioning"
+    #     )
+    #     self.feature_extractor = ViTImageProcessor.from_pretrained(
+    #         "nlpconnect/vit-gpt2-image-captioning"
+    #     )
+    #     self.tokenizer = AutoTokenizer.from_pretrained(
+    #         "nlpconnect/vit-gpt2-image-captioning"
+    #     )
+
+    #     # device = torch.device("cuda" if torch.cuda.is_available() else "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
+    #     self.gen_kwargs = {"max_length": self.max_length, "num_beams": self.num_beams}
 
 
     def predict_step(self,image_paths):
@@ -45,6 +50,39 @@ 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 predict_image_caption_gemini(self,img):
+        prompt = "Describe the main focus of this image in detail."
+        response = self.model.generate_content([prompt, img], stream=True)
+        response.resolve()
+        print("Derived data",response.text)
+        return response.text
+
     def get_args(self):
         parser = argparse.ArgumentParser()
         parser.add_argument( "-i",