m7mdal7aj commited on
Commit
15d3f2d
1 Parent(s): c200cd2

Update my_model/captioner/image_captioning.py

Browse files
my_model/captioner/image_captioning.py CHANGED
@@ -66,7 +66,10 @@ class ImageCaptioningModel:
66
 
67
  def generate_caption(self, image_path):
68
 
69
- image = Image.open(image_path)
 
 
 
70
  image = self.resize_image(image)
71
  inputs = self.processor(image, self.prompt, return_tensors="pt").to("cuda", self.torch_dtype)
72
  outputs = self.model.generate(**inputs, min_length=self.min_length, max_new_tokens=self.max_new_tokens)
 
66
 
67
  def generate_caption(self, image_path):
68
 
69
+ if isinstance(image, str) or isinstance(image, io.IOBase):
70
+ # If it's a file path or file-like object, open it as a PIL Image
71
+ image = Image.open(image_path)
72
+
73
  image = self.resize_image(image)
74
  inputs = self.processor(image, self.prompt, return_tensors="pt").to("cuda", self.torch_dtype)
75
  outputs = self.model.generate(**inputs, min_length=self.min_length, max_new_tokens=self.max_new_tokens)