ImageCaption / app.py
Thunder-rk's picture
Create app.py
a9156aa verified
raw
history blame
998 Bytes
import gradio as gr
from transformers import BlipProcessor, BlipForConditionalGeneration
from PIL import Image
processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base")
def generate_caption(image):
# Now directly using the PIL Image object
inputs = processor(images=image, return_tensors="pt")
outputs = model.generate(**inputs)
caption = processor.decode(outputs[0], skip_special_tokens=True)
return caption
def caption_image(image):
"""
Takes a PIL Image input and returns a caption.
"""
try:
caption = generate_caption(image)
return caption
except Exception as e:
return f"An error occurred: {str(e)}"
iface = gr.Interface(
fn=caption_image,
inputs=gr.Image(type="pil"),
outputs="text",
title="Image Captioning with BLIP",
description="Upload an image to generate a caption."
)
iface.launch()