| import gradio as gr
|
| from transformers import BlipProcessor, BlipForConditionalGeneration
|
| from PIL import Image
|
| import torch
|
| from transformers import GPT2Tokenizer, GPT2LMHeadModel
|
|
|
| device = 'cuda' if torch.cuda.is_available() else 'cpu'
|
|
|
| processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
|
| model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base").to(device)
|
|
|
| tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
|
| gpt2_model = GPT2LMHeadModel.from_pretrained("gpt2").to(device)
|
|
|
| def generate_paragraph(image):
|
| if image.mode != 'RGB':
|
| image = image.convert('RGB')
|
| inputs = processor(images=image, return_tensors="pt").to(device)
|
| output_ids = model.generate(**inputs, max_length=50)
|
| caption = processor.decode(output_ids[0], skip_special_tokens=True)
|
|
|
| prompt = f"Write a detailed paragraph about this image: {caption}\n\nDetails:"
|
| tokens = tokenizer.encode(prompt, return_tensors='pt').to(device)
|
| outputs = gpt2_model.generate(tokens, max_length=150, num_beams=5, no_repeat_ngram_size=2, early_stopping=True, pad_token_id=tokenizer.eos_token_id)
|
| paragraph = tokenizer.decode(outputs[0], skip_special_tokens=True)
|
|
|
|
|
| if paragraph.lower().startswith(prompt.lower()):
|
| paragraph = paragraph[len(prompt):].strip()
|
|
|
| return paragraph
|
|
|
| iface = gr.Interface(
|
| fn=generate_paragraph,
|
| inputs=gr.Image(type="pil"),
|
| outputs="textbox",
|
| title="Image Paragraph Description Generator",
|
| description="Upload an image to get a detailed paragraph description generated."
|
| )
|
|
|
| iface.launch()
|
|
|
