import gradio as gr
from transformers import pipeline
from PIL import Image

# 加载预训练的图像文本描述模型
image_to_text = pipeline("image-to-text", model="nlpconnect/vit-gpt2-image-captioning")

# 定义核心生成函数
def generate_caption(image):
    # 处理图像并生成描述
    results = image_to_text(image)
    # 提取生成的文本内容，格式化输出
    return f"图像描述：{results[0]['generated_text']}"

# 搭建Gradio交互式界面
with gr.Blocks(title="图像文本描述生成工具") as demo:
    gr.Markdown("# 图像文本描述生成工具")
    gr.Markdown("上传任意图片，将自动生成对应的英文文本描述~")
    # 图像上传组件
    image_input = gr.Image(type="pil", label="上传图片")
    # 结果输出组件
    text_output = gr.Textbox(label="生成的文本描述")
    # 按钮绑定函数，点击触发生成
    gr.Button("生成描述").click(fn=generate_caption, inputs=image_input, outputs=text_output)

if __name__ == "__main__":
    demo.launch()