Possible same

utils/caption_utils.py

@@ -2,27 +2,18 @@ from PIL import Image
 import io
 import torch
 from transformers import BlipProcessor, BlipForConditionalGeneration
-from utils.image_utils import load_image
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-
 class ImageCaptioning:
 
     def __init__(self):
         self.processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
         self.model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base").to(device)
 
-    def get_caption(self, image):
-
-        # Preprocessing the Image
-        img = Image.open(io.BytesIO(image))
-        img = self.processor(img, return_tensors="pt").to(device)
-
-        # Generating captions
-        output = self.model.generate(**img)
-
-        # decode the output
+    def get_caption(self, image_bytes):
+        img = Image.open(io.BytesIO(image_bytes))
+        img_tensors = self.processor(img, return_tensors="pt").to(device)
+        output = self.model.generate(**img_tensors)
         caption = self.processor.batch_decode(output, skip_special_tokens=True)[0]
-
         return caption