Update app.py

app.py

@@ -1,64 +1,87 @@
-# import part
 import streamlit as st
 from transformers import pipeline
-import torch
+from datasets import load_dataset
+from PIL import Image
+import numpy as np
+from collections import Counter
 
-# function part
-# img2text
-def img2text(url):
-    image_to_text_model = pipeline("image-to-text", model="Salesforce/blip-image-captioning-base")
-    text = image_to_text_model(url)[0]["generated_text"]
-    return text
+# 设置页面标题
+st.set_page_config(page_title="🏠 Airbnb Design Advisor", layout="wide")
 
-# text2story
-def text2story(text):
-    pipe = pipeline("text-generation", model="gpt2")
-    story_text = pipe(text, max_length=100)[0]['generated_text']
-    return story_text
+# 缓存模型和数据集
+@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")
+    }
 
-# text2audio
-def text2audio(story_text):
-    pipe = pipeline("text-to-speech", model="facebook/mms-tts-eng")
-    audio_data = pipe(story_text)
-    return audio_data
+@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_]
 
-# main part
+# 主界面
+st.title("AI-Powered Airbnb Design Advisor")
 
-st.set_page_config(page_title="Your Image to Audio Story",
-                   page_icon="🦜")
-st.header("Turn Your Image to Audio Story")
-uploaded_file = st.file_uploader("Select an Image...")#, type=["jpg", "jpeg", "png"])
-
-
-if uploaded_file is not None:
-    print(uploaded_file)
-    bytes_data = uploaded_file.getvalue()
-    with open(uploaded_file.name, "wb") as file:
-        file.write(bytes_data)
-    st.image(uploaded_file, caption="Uploaded Image", use_column_width=True)
+# 侧边栏控制
+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()
     
-
-    #Stage 1: Image to Text
-    st.text('Processing img2text...')
-    scenario = img2text(uploaded_file.name)
-    st.write(scenario)
-
-
-    #Stage 2: Text to Story
-    st.text('Generating a story...')
-    story = text2story(scenario)
-    st.write(story)
-
-    #Stage 3: Story to Audio data
-    st.text('Generating audio data...')
-    audio_data =text2audio(story)
-
-    # Play button
-    if st.button("Play Audio"):
-        st.audio(audio_data['audio'],
-                   format="audio/wav",
-                   start_time=0,
-                   sample_rate = audio_data['sampling_rate'])
-        # st.audio("kids_playing_audio.wav")
+    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']
+        
+        # 显示结果
+        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.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 col2:
+            st.subheader("Professional Advice")
+            st.write(advice)
+        
+        st.subheader("Example Images")
+        st.image(examples['image'].tolist(), width=300)