| --- |
| license: mit |
| --- |
| Part of Advanced NLP Project for Team Shrine - Adnan Qidwai, Harshit Karwal and Shrikara Arun. |
| CleanCaption is an image captioning model that forget an object from the image when generating the caption. It is a finetuned version of `microsoft/Florence-2-large-ft`. |
|
|
| Usage: |
| ```python |
| from transformers import AutoProcessor, AutoModelForCausalLM |
| from PIL import Image |
| import torch |
| |
| device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu" |
| |
| processor = AutoProcessor.from_pretrained("microsoft/Florence-2-large-ft", trust_remote_code=True) |
| model = AutoModelForCausalLM.from_pretrained( |
| "sudokara/CleanCaption", |
| trust_remote_code=True |
| ).eval().to(device) |
| |
| def forget(prompt, image_path): |
| image = Image.open(image_path).convert("RGB") |
| prompt = f"Forget from caption: {str(prompt)}".strip(' :') |
| inputs = processor(text=prompt, images=image, return_tensors="pt").to(device) |
| generated_ids = model.generate( |
| input_ids=inputs["input_ids"], |
| pixel_values=inputs["pixel_values"], |
| max_new_tokens=1024, |
| do_sample=True, |
| num_beams=3, |
| ) |
| return processor.decode(generated_ids[0]).replace('<s>', '').replace('</s>', '') |
| |
| image_path = "image.png" |
| print(forget(image_path = image_path, prompt = "water")) |
| ``` |
|
|
