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
|
|
|
|
|
|
|
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)
|