From 81847c24836ac15044eca47c0024ef72ccb09897 Mon Sep 17 00:00:00 2001 From: craig_wright156 Date: Fri, 20 Oct 2023 19:41:30 +0100 Subject: [PATCH] imageCaptioning fix --- nodes/ImageCaptioningNode.py | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/nodes/ImageCaptioningNode.py b/nodes/ImageCaptioningNode.py index 726d825..bceea0d 100644 --- a/nodes/ImageCaptioningNode.py +++ b/nodes/ImageCaptioningNode.py @@ -21,13 +21,12 @@ class ImageCaptioningNode: self.model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-large").to("cuda") def caption(self, image): - image = image.numpy() - if image.ndim == 4: # image has batch dimension - image = image[0] # take first image in batch - image = Image.fromarray((image * 255).astype(np.uint8).transpose(1, 2, 0)) - + + i = 255. * image[0].cpu().numpy() + img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8)) + # Perform unconditional image captioning - inputs = self.processor(image, return_tensors="pt").to("cuda") + inputs = self.processor(img, return_tensors="pt").to("cuda") out = self.model.generate(**inputs) caption = self.processor.decode(out[0], skip_special_tokens=True)