Update app.py

app.py

@@ -1,14 +1,111 @@
-# 测试文本模型
-# import part
-import streamlit as st
-from transformers import pipeline
+import gradio as gr
 import torch
-# pipe = pipeline("text-classification", model="distilbert-base-uncased")
-# print(pipe("This is a test sentence."))
-
-# # 测试视觉模型
-# from PIL import Image
-# import requests
-# pipe = pipeline("image-classification", model="google/vit-base-patch16-224")
-# url = "http://images.cocodataset.org/val2017/000000039769.jpg"
-# print(pipe(Image.open(requests.get(url, stream=True).raw)))
+from transformers import pipeline
+from datasets import load_dataset
+from PIL import Image
+import numpy as np
+from collections import Counter
+import functools
+
+# 使用标准库的缓存装饰器替代Gradio缓存
+@functools.lru_cache(maxsize=None)
+def load_models():
+    return {
+        "detector": pipeline(
+            "object-detection", 
+            model="facebook/detr-resnet-50",
+            device=0 if torch.cuda.is_available() else -1
+        ),
+        "generator": pipeline(
+            "text2text-generation",
+            model="google/flan-t5-base",  # 改用基础版降低资源需求
+            device=0 if torch.cuda.is_available() else -1
+        )
+    }
+
+# 数据集加载函数（移除Gradio缓存）
+def load_dataset_data():
+    ds = load_dataset("AntZet/home_decoration_objects_images")
+    return ds['train'].to_pandas()
+
+# 颜色分析函数保持不变
+def get_dominant_colors(img, n_colors=3):
+    arr = np.array(img.resize((100,100)))
+    pixels = arr.reshape(-1,3)
+    from sklearn.cluster import KMeans
+    kmeans = KMeans(n_clusters=n_colors)
+    kmeans.fit(pixels)
+    return [f"#{int(c[0]):02x}{int(c[1]):02x}{int(c[2]):02x}" for c in kmeans.cluster_centers_]
+
+# 核心处理函数
+def generate_recommendation(target_style):
+    try:
+        models = load_models()
+        df = load_dataset_data()
+        
+        style_df = df[df['style'] == target_style.lower()]
+        if len(style_df) < 3:
+            return f"⚠️ Not enough samples for {target_style} style"
+        
+        sample_images = style_df.sample(5)['image']
+        
+        all_objects = []
+        color_palette = []
+        
+        for img in sample_images:
+            detected = models["detector"](img)
+            all_objects += [obj['label'] for obj in detected if obj['score'] > 0.9]
+            color_palette += get_dominant_colors(img)
+        
+        top_objects = Counter(all_objects).most_common(3)
+        top_colors = Counter(color_palette).most_common(3)
+        
+        prompt = f"""Create interior design recommendations for {target_style} style:
+        Key objects: {[o[0] for o in top_objects]}
+        Color palette: {[c[0] for c in top_colors]}
+        Include: 3 essentials, 2 budget tips, common mistakes"""
+        
+        advice = models["generator"](prompt, max_length=300)[0]['generated_text']
+        
+        output = f"## 🎨 {target_style.title()} Style Guide\n\n"
+        output += "### 🪑 Key Objects\n" + "\n".join(
+            [f"- {o[0]} ({o[1]}x)" for o in top_objects]) + "\n\n"
+        output += "### 🎨 Colors\n" + "\n".join(
+            [f"<span style='color:{c[0]};'>■</span> {c[0]}" for c in top_colors]) + "\n\n"
+        output += "### 💡 Advice\n" + advice.replace(". ", ".\n")
+        
+        return output
+    
+    except Exception as e:
+        return f"❌ Error: {str(e)}"
+
+# Gradio界面保持不变
+with gr.Blocks(title="Design Assistant") as demo:
+    gr.Markdown("# 🏡 AI Design Advisor")
+    
+    with gr.Row():
+        style_input = gr.Dropdown(
+            label="Select Style",
+            choices=["Industrial", "Scandinavian", "Bohemian", "Modern"],
+            value="Industrial"
+        )
+    
+    submit_btn = gr.Button("Generate Plan", variant="primary")
+    
+    with gr.Row():
+        output = gr.Markdown()
+        gallery = gr.Gallery(
+            label="Examples",
+            object_fit="contain",
+            height="300px"
+        )
+    
+    def update_gallery(style):
+        df = load_dataset_data()
+        return df[df['style'] == style.lower()].sample(3)['image'].tolist()
+    
+    style_input.change(update_gallery, inputs=style_input, outputs=gallery)
+    submit_btn.click(generate_recommendation, inputs=style_input, outputs=output)
+
+if __name__ == "__main__":
+    demo.launch()