Spaces:
Sleeping
Sleeping
import gradio as gr | |
from transformers import pipeline | |
# Load the image-to-text pipeline | |
image_to_text_pipelines = { | |
"Salesforce/blip-image-captioning-base": pipeline("image-to-text", model="Salesforce/blip-image-captioning-base"), | |
# Add more models if needed | |
} | |
def generate_caption(input_image, model_name="Salesforce/blip-image-captioning-base"): | |
# Generate caption for the input image using the selected model | |
image_to_text_pipeline = image_to_text_pipelines[model_name] | |
caption = image_to_text_pipeline(input_image)[0]['generated_text'] | |
return caption | |
# Interface for launching the model | |
interface = gr.Interface( | |
fn=generate_caption, | |
inputs=gr.Image(type='pil', label="Input Image"), | |
outputs="text", | |
title="Image Captioning Model", | |
description="This model generates captions for images.", | |
theme="default", | |
) | |
# Launch the interface | |
interface.launch() |