import gradio as gr
from transformers import pipeline
from PIL import Image
import jieba.analyse

# 初始化多语言描述管道（英文生成+中英翻译）
image_to_text = pipeline("image-to-text", model="IDEA-CCNL/Taiyi-BLIP-750M-Chinese")

def analyze_image(image):
    # 生成英文描述
    zh_result = image_to_text(image)[0]['generated_text']
    print(zh_result)
    
    # 提取关键词
    keywords = jieba.analyse.textrank(
        zh_result,
        topK=5,
        allowPOS=('n', 'vn', 'v')
    )
    
    # 格式化输出
    return f"图片描述：{zh_result}\n关键字：{', '.join(keywords)}\n"  # 第三行为空行保持格式

# 创建 Gradio 界面
demo = gr.Interface(
    fn=analyze_image,
    inputs=gr.Image(type="pil", label="上传商品图片"),
    outputs=gr.Textbox(label="分析结果", lines=3),
    title="",
    description="上传商品图片，自动生成中文描述并提取关键词",
    allow_flagging="never"
)

# 运行应用
if __name__ == "__main__":
    demo.launch()