Update app.py

app.py

@@ -3,85 +3,84 @@ 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="🏠 装修风格分析器", layout="wide")
+st.title("AI 装修风格匹配工具")
 
-        
-# 初始化模型 (缓存)
+# 缓存模型（移除了物体检测）
 @st.cache_resource
 def load_models():
     return {
-        "detector": pipeline("object-detection", model="facebook/detr-resnet-50"),
-        "style_classifier":pipeline("image-classification", model="playrobin/furniture-styles"
+        "style_classifier": pipeline(
+            "image-classification", 
+            model="dima806/interior_design_style_classification"
         ),
-        "advisor": pipeline("text2text-generation", model="google/flan-t5-base")
+        "advisor": pipeline("text2text-generation", model="google/flan-t5-small")
     }
 
-# 加载数据集 (缓存)
-@st.cache_data
-def load_style_examples():
-    dataset = load_dataset("AntZet/home_decoration_objects_images")
-    return dataset['train'].to_pandas()
-
-@st.cache_data
-def load_style_examples():
-    try:
-        dataset = load_dataset("AntZet/home_decoration_objects_images", streaming=True)
-        return dataset['train'].to_pandas()
-    except:
-        return pd.DataFrame(columns=['style', 'image'])  # 返回空数据集作为fallback
+# 颜色分析函数（替代物体检测）
+def analyze_image(img):
+    # 简化的视觉分析：仅提取颜色
+    img = img.resize((50,50))
+    arr = np.array(img)
+    pixels = arr.reshape(-1,3)
+    
+    # 使用简化版颜色分析（避免sklearn依赖）
+    unique_colors = np.unique(pixels, axis=0)
+    main_colors = unique_colors[:3]  # 取前3种主要颜色
+    return [f"#{r:02x}{g:02x}{b:02x}" for r,g,b in main_colors]
 
-        
-# 主函数
 def main():
-    st.title("🏠 AI 装修风格匹配器")
-    
     uploaded_img = st.file_uploader("上传房间照片", type=["jpg", "png"])
     
     if uploaded_img:
-        img = Image.open(uploaded_img)
         models = load_models()
-        df = load_style_examples()
+        img = Image.open(uploaded_img)
+        
+        with st.spinner("正在分析..."):
+            # 1. 风格分类
+            style_result = models["style_classifier"](img)
+            main_style = style_result[0]['label']
+            
+            # 2. 视觉分析（颜色替代物体检测）
+            colors = analyze_image(img)
+            
+            # 3. 从数据集找案例
+            try:
+                dataset = load_dataset("AntZet/home_decoration_objects_images", streaming=True)
+                examples = [ex['image'] for ex in dataset['train'] 
+                           if ex['style'] == main_style][:3]
+            except:
+                examples = []
+            
+            # 4. 生成建议
+            prompt = f"""基于{main_style}风格，给出3条装修建议：
+            - 主色调: {colors}
+            - 避免: 与风格冲突的元素
+            - 预算: 中等成本方案"""
+            advice = models["advisor"](prompt, max_length=200)[0]['generated_text']
         
+        # 显示结果
         col1, col2 = st.columns(2)
         
         with col1:
             st.image(img, width=300)
+            st.success(f"识别风格: {main_style}")
             
-            # 风格分类
-            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)))
-
+            st.subheader("主要色调")
+            for color in colors:
+                st.markdown(f"<div style='background:{color}; height:30px'></div>", 
+                           unsafe_allow_html=True)
+        
         with col2:
-            # 从数据集中找匹配案例
-            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.subheader("风格建议")
             st.write(advice)
+            
+            if examples:
+                st.subheader("参考案例")
+                st.image(examples, width=150)
 
 if __name__ == "__main__":
     main()