Spaces:
Sleeping
Sleeping
| # from agentlego.tools import BaseTool | |
| # from PIL import Image | |
| # import torch | |
| # class ImageDescriptionTool(BaseTool): | |
| # default_desc = 'Uses a pretrained BLIP model to generate descriptions for images.' | |
| # def __init__(self): | |
| # super().__init__() | |
| # # Load models inside the class initialization | |
| # from transformers import AutoProcessor, AutoModelForImageTextToText | |
| # MODEL_ID = "Salesforce/blip-image-captioning-base" | |
| # self.processor = AutoProcessor.from_pretrained(MODEL_ID) | |
| # self.model = AutoModelForImageTextToText.from_pretrained(MODEL_ID) | |
| # # Set up device and generation parameters | |
| # self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| # self.model.to(self.device) | |
| # self.max_length = 256 | |
| # self.num_beams = 5 | |
| # self.gen_kwargs = { | |
| # "max_length": self.max_length, | |
| # "num_beams": self.num_beams, | |
| # "early_stopping": True | |
| # } | |
| # def apply(self, image_path: str) -> str: | |
| # try: | |
| # # Open the image | |
| # image = Image.open(image_path) | |
| # if image.mode != "RGB": | |
| # image = image.convert(mode="RGB") | |
| # # Preprocess image | |
| # inputs = self.processor(images=image, return_tensors="pt").to(self.device) | |
| # # Generate caption | |
| # with torch.no_grad(): | |
| # output_ids = self.model.generate(**inputs, **self.gen_kwargs) | |
| # # Decode prediction | |
| # caption = self.processor.batch_decode(output_ids, skip_special_tokens=True)[0].strip() | |
| # return f"Description: **{caption}** (generated with BLIP base model)" | |
| # except Exception as e: | |
| # return f"Error during image description: {str(e)}" | |
| from agentlego.tools import BaseTool | |
| from PIL import Image | |
| import torch | |
| class ImageDescriptionTool(BaseTool): | |
| default_desc = 'Uses a pretrained VIT-GPT2 model to generate descriptions for images.' | |
| def __init__(self): | |
| super().__init__() | |
| # Load models inside the class initialization | |
| from transformers import VisionEncoderDecoderModel, ViTImageProcessor, AutoTokenizer | |
| self.model = VisionEncoderDecoderModel.from_pretrained("nlpconnect/vit-gpt2-image-captioning") | |
| self.feature_extractor = ViTImageProcessor.from_pretrained("nlpconnect/vit-gpt2-image-captioning") | |
| self.tokenizer = AutoTokenizer.from_pretrained("nlpconnect/vit-gpt2-image-captioning") | |
| # Set up device and generation parameters | |
| self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| self.model.to(self.device) | |
| self.max_length = 16 | |
| self.num_beams = 4 | |
| self.gen_kwargs = {"max_length": self.max_length} # no num_beams = greedy decoding | |
| def apply(self, image_path: str) -> str: | |
| try: | |
| # Open the image | |
| image = Image.open(image_path) | |
| if image.mode != "RGB": | |
| image = image.convert(mode="RGB") | |
| # Preprocess image | |
| pixel_values = self.feature_extractor(images=[image], return_tensors="pt").pixel_values | |
| pixel_values = pixel_values.to(self.device) | |
| # Generate caption | |
| with torch.no_grad(): | |
| output_ids = self.model.generate(pixel_values, **self.gen_kwargs) | |
| # Decode prediction | |
| pred = self.tokenizer.decode(output_ids[0], skip_special_tokens=True) | |
| pred = pred.strip() | |
| return f"Description: **{pred}** (generated with VIT-GPT2 model)" | |
| except Exception as e: | |
| return f"Error during image description: {str(e)}" | |