| 299 | |
| 300 | |
| 301 | class ImageCaptioning: |
| 302 | def __init__(self, device): |
| 303 | print(f"Initializing ImageCaptioning to {device}") |
| 304 | self.device = device |
| 305 | self.torch_dtype = torch.float16 if 'cuda' in device else torch.float32 |
| 306 | self.processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base") |
| 307 | self.model = BlipForConditionalGeneration.from_pretrained( |
| 308 | "Salesforce/blip-image-captioning-base", torch_dtype=self.torch_dtype).to(self.device) |
| 309 | |
| 310 | @prompts(name="Get Photo Description", |
| 311 | description="useful when you want to know what is inside the photo. receives image_path as input. " |
| 312 | "The input to this tool should be a string, representing the image_path. ") |
| 313 | def inference(self, image_path): |
| 314 | inputs = self.processor(Image.open(image_path), return_tensors="pt").to(self.device, self.torch_dtype) |
| 315 | out = self.model.generate(**inputs) |
| 316 | captions = self.processor.decode(out[0], skip_special_tokens=True) |
| 317 | print(f"\nProcessed ImageCaptioning, Input Image: {image_path}, Output Text: {captions}") |
| 318 | return captions |
| 319 | |
| 320 | |
| 321 | class Image2Canny: |
nothing calls this directly
no outgoing calls
no test coverage detected