Update app.py

app.py

@@ -3,85 +3,74 @@ from transformers import pipeline
 from datasets import load_dataset
 from PIL import Image
 import numpy as np
-from collections import Counter
 
-# 设置页面标题
-st.set_page_config(page_title="🏠 Airbnb Design Advisor", layout="wide")
-
-# 缓存模型和数据集
+# 初始化模型 (缓存)
 @st.cache_resource
 def load_models():
     return {
         "detector": pipeline("object-detection", model="facebook/detr-resnet-50"),
-        "generator": pipeline("text2text-generation", model="google/flan-t5-base")
+        "style_classifier":pipeline("image-classification", model="playrobin/furniture-styles"
+        ),
+        "advisor": pipeline("text2text-generation", model="google/flan-t5-base")
     }
 
+# 加载数据集 (缓存)
 @st.cache_data
-def load_data():
-    return load_dataset("AntZet/home_decoration_objects_images")['train'].to_pandas()
-
-# 颜色提取函数
-def get_colors(img, n=3):
-    arr = np.array(img.resize((50,50)))
-    from sklearn.cluster import KMeans
-    return [f"#{int(c[0]):02x}{int(c[1]):02x}{int(c[2]):02x}" 
-            for c in KMeans(n_clusters=n).fit(arr.reshape(-1,3)).cluster_centers_]
-
-# 主界面
-st.title("AI-Powered Airbnb Design Advisor")
+def load_style_examples():
+    dataset = load_dataset("AntZet/home_decoration_objects_images")
+    return dataset['train'].to_pandas()
 
-# 侧边栏控制
-with st.sidebar:
-    st.header("Settings")
-    style = st.selectbox(
-        "Select Style",
-        ["industrial", "scandinavian", "bohemian", "modern"],
-        index=0
-    )
-    analyze_btn = st.button("Analyze")
-
-# 主内容区
-if analyze_btn:
-    models = load_models()
-    df = load_data()
+# 主函数
+def main():
+    st.title("🏠 AI 装修风格匹配器")
     
-    with st.spinner("Generating recommendations..."):
-        # 获取风格示例
-        examples = df[df['style'] == style].sample(3)
-        
-        # 分析对象和颜色
-        objects = []
-        colors = []
-        for img in examples['image']:
-            detected = models["detector"](img)
-            objects += [obj['label'] for obj in detected if obj['score'] > 0.8]
-            colors += get_colors(img)
-        
-        # 生成建议
-        prompt = f"""Create {style} style decoration tips for Airbnb with:
-        - Key objects: {Counter(objects).most_common(3)}
-        - Color palette: {Counter(colors).most_common(3)}
-        Include: 3 essentials, 2 budget tips"""
-        
-        advice = models["generator"](prompt, max_length=300)[0]['generated_text']
+    uploaded_img = st.file_uploader("上传房间照片", type=["jpg", "png"])
+    
+    if uploaded_img:
+        img = Image.open(uploaded_img)
+        models = load_models()
+        df = load_style_examples()
         
-        # 显示结果
         col1, col2 = st.columns(2)
         
         with col1:
-            st.subheader("Key Elements")
-            for obj, count in Counter(objects).most_common(3):
-                st.markdown(f"- {obj} (appears in {count} samples)")
+            st.image(img, width=300)
             
-            st.subheader("Color Palette")
-            for color in Counter(colors).most_common(3):
-                st.markdown(f"<div style='background:{color[0]}; width:100%; height:30px'></div>", 
-                           unsafe_allow_html=True)
-                st.caption(color[0])
-        
+            # 风格分类
+            with st.spinner("分析房间风格..."):
+                style_result = models["style_classifier"](img)
+                main_style = style_result[0]['label']
+                confidence = style_result[0]['score']
+            
+            st.success(f"检测风格: {main_style} (置信度: {confidence:.0%})")
+            
+            # 物体检测
+            with st.spinner("识别家具物品..."):
+                objects = models["detector"](img)
+                top_objects = [obj['label'] for obj in objects if obj['score'] > 0.7]
+            
+            st.subheader("检测到的主要物品")
+            st.write(", ".join(set(top_objects)))
+
         with col2:
-            st.subheader("Professional Advice")
+            # 从数据集中找匹配案例
+            style_samples = df[df['style'] == main_style].sample(3)
+            
+            st.subheader(f"{main_style} 风格案例")
+            st.image(style_samples['image'].tolist(), width=150)
+            
+            # 生成建议
+            with st.spinner("生成装修建议..."):
+                prompt = f"""根据以下条件生成装修建议：
+                - 当前风格: {main_style}
+                - 现有物品: {top_objects[:5]}
+                - 目标风格: {main_style}
+                提供3条具体改进建议"""
+                
+                advice = models["advisor"](prompt, max_length=300)[0]['generated_text']
+            
+            st.subheader("专业建议")
             st.write(advice)
-        
-        st.subheader("Example Images")
-        st.image(examples['image'].tolist(), width=300)
+
+if __name__ == "__main__":
+    main()