Update app.py

app.py

@@ -1,36 +1,126 @@
 import gradio as gr
-# from transformers import pipeline
-
-# 初始化文本生成管道
-try:
-    generator = pipeline("text-generation", model="distilgpt2")
-except Exception as e:
-    print(f"模型加载失败: {e}")
-    generator = None
-
-def generate_advice(style):
-    if not generator:
-        return "系统初始化失败，请检查transformers库是否安装", ""
-    
-    prompt = f"如何将Airbnb房间装修成{style}风格？请给出3条具体建议"
-    try:
-        result = generator(prompt, max_length=200)
-        return result[0]["generated_text"], ""
-    except Exception as e:
-        return f"生成建议时出错: {str(e)}", ""
-
-# 创建简单界面
-with gr.Blocks() as demo:
-    gr.Markdown("## 🏡 Airbnb装修助手")
-    style_input = gr.Textbox(label="输入想要的风格")
-    submit_btn = gr.Button("生成建议")
-    advice_output = gr.Textbox(label="装修建议")
+from transformers import pipeline
+from PIL import Image
+import requests
+from io import BytesIO
+import numpy as np
+
+# 初始化两个pipeline
+item_classifier = pipeline(
+    "image-classification",
+    model="google/vit-base-patch16-224"
+)
+
+scene_analyzer = pipeline(
+    "image-to-text",
+    model="nlpconnect/vit-gpt2-image-captioning"
+)
+
+# 风格映射字典 (可根据实际需求扩展)
+STYLE_MAPPING = {
+    "modern": "现代风格",
+    "contemporary": "当代风格",
+    "minimalist": "极简风格",
+    "industrial": "工业风格",
+    "scandinavian": "北欧风格",
+    "bohemian": "波西米亚风格",
+    "rustic": "乡村风格",
+    "traditional": "传统风格",
+    "coastal": "海岸风格",
+    "mid-century": "中世纪现代风格"
+}
+
+def determine_overall_style(items, description):
+    """
+    根据物品分类和场景描述确定整体风格
+    """
+    # 从物品分类中提取关键词
+    item_keywords = " ".join([item['label'] for item in items])
+    
+    # 合并所有文本信息
+    combined_text = f"{item_keywords} {description}".lower()
+    
+    # 简单的关键词匹配确定风格
+    detected_styles = []
+    for style_en, style_cn in STYLE_MAPPING.items():
+        if style_en in combined_text:
+            detected_styles.append(style_cn)
+    
+    # 如果没有匹配到任何风格，返回一个默认值
+    if not detected_styles:
+        return "混合风格"
+    
+    # 返回匹配到的所有风格
+    return "、".join(detected_styles)
+
+def analyze_room_style(image):
+    """
+    分析房间风格的主函数
+    """
+    # 如果是URL，下载图片
+    if isinstance(image, str) and image.startswith(('http://', 'https://')):
+        response = requests.get(image)
+        image = Image.open(BytesIO(response.content))
+    
+    # 物品级别分析 (取前5个最相关的物品)
+    item_results = item_classifier(image, top_k=5)
+    
+    # 场景级别分析
+    scene_description = scene_analyzer(image)[0]['generated_text']
+    
+    # 综合判断风格
+    style = determine_overall_style(item_results, scene_description)
+    
+    # 格式化结果
+    items_formatted = "\n".join([f"- {item['label']} ({item['score']:.2f})" for item in item_results])
+    
+    return {
+        "detected_items": items_formatted,
+        "scene_description": scene_description,
+        "predicted_style": style
+    }
+
+def predict_style(image):
+    """
+    Gradio接口使用的预测函数
+    """
+    result = analyze_room_style(image)
+    
+    output = f"""🏠 预测风格: {result['predicted_style']}
+    
+📝 场景描述: {result['scene_description']}
+    
+🛋️ 检测到的物品:
+{result['detected_items']}
+"""
+    return output
+
+# 创建Gradio界面
+with gr.Blocks(title="Airbnb房屋风格识别") as demo:
+    gr.Markdown("# 🏡 Airbnb房屋风格识别")
+    gr.Markdown("上传您的房间照片，AI将分析您的房屋装饰风格")
+    
+    with gr.Row():
+        with gr.Column():
+            image_input = gr.Image(type="filepath", label="上传房间照片")
+            submit_btn = gr.Button("分析风格")
+        with gr.Column():
+            output = gr.Textbox(label="分析结果", lines=10)
+    
+    examples = gr.Examples(
+        examples=[
+            ["https://example.com/room1.jpg"],  # 替换为实际示例图片URL
+            ["https://example.com/room2.jpg"]   # 替换为实际示例图片URL
+        ],
+        inputs=image_input
+    )
     
     submit_btn.click(
-        fn=generate_advice,
-        inputs=style_input,
-        outputs=advice_output
+        fn=predict_style,
+        inputs=image_input,
+        outputs=output
     )
 
+# 启动应用
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch(server_name="0.0.0.0", server_port=7860)