Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,40 +1,161 @@
-import altair as alt
+import streamlit as st
+import os, pickle
+import torch
 import numpy as np
 import pandas as pd
-import streamlit as st
+from PIL import Image as PILImage
+from sentence_transformers import SentenceTransformer
+from transformers import AutoProcessor, Gemma3ForConditionalGeneration, CLIPProcessor, CLIPModel
+import faiss
+
+from huggingface_hub import login
+# Paste your token inside the quotes
+login(st.secrets["huggingface"]["token"])
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+print(device)
+
+# Load assets
+@st.cache_resource
+def load_assets(asset_dir="Assets"):
+    with open(os.path.join(asset_dir, "image_embeddings.pkl"), "rb") as f:
+        image_embeddings = pickle.load(f)
+    with open(os.path.join(asset_dir, "text_embeddings.pkl"), "rb") as f:
+        text_embeddings = pickle.load(f)
+    with open(os.path.join(asset_dir, "product_ids.pkl"), "rb") as f:
+        ids = pickle.load(f)
+    combined_vectors = np.load(os.path.join(asset_dir, "combined_vectors.npy"))
+    faiss_index = faiss.read_index(os.path.join(asset_dir, "faiss_index.index"))
+    df = pd.read_pickle(os.path.join(asset_dir, "product_metadata_df.pkl"))
+    with open(os.path.join(asset_dir, "user_history.pkl"), "rb") as f:
+        user_history = pickle.load(f)
+    with open(os.path.join(asset_dir, "trend_string.pkl"), "rb") as f:
+        trend_string = pickle.load(f)
+    return image_embeddings, text_embeddings, ids, combined_vectors, faiss_index, df, user_history, trend_string
+
+# Image + text search
+def search_similar(image=None, text=None, top_k=5):
+    img_vec = np.zeros(768)
+    txt_vec = np.zeros(384)
+    if image:
+        inputs = clip_processor(images=image, return_tensors="pt").to(device)
+        with torch.no_grad():
+            img_vec = clip_model.get_image_features(**inputs).cpu().numpy()[0]
+    if text:
+        txt_vec = text_model.encode(text)
+    combined = np.concatenate([img_vec, txt_vec]).astype("float32")
+    D, I = faiss_index.search(np.array([combined]), top_k)
+    return [ids[i] for i in I[0]]
+
+# Outfit suggestions
+def generate_outfit_gemma(img, row, username, suggestions=5):
+    brands, styles, desc = summarize_user_preferences(username)
+    messages = [{
+        "role": "system",
+        "content": [{"type": "text", "text": "You are a highly experienced fashion stylist and personal shopper."}]
+    }, {
+        "role": "user",
+        "content": [
+            {"type": "image", "image": img.convert("RGB")},
+            {"type": "text", "text": f"""
+Suggest {suggestions} stylish outfit items that complement this item:
+
+**Product**:
+Name: {row['product_name']}
+Brand: {row['brand']}
+Style: {row['style_attributes']}
+Description: {row['description']}
+Price: ₹{row['selling_price']}
+
+**User Likes**:
+Brands: {brands}
+Styles: {styles}
+Liked Items: {desc}
+
+**Trends**:
+{trend_string}
+
+Output in bullet list with name + explanation.
+            """}
+        ]
+    }]
+    prompt = gemma_processor.apply_chat_template(messages, add_generation_prompt=True, tokenize=False)
+    tokenized = gemma_processor(text=prompt, return_tensors="pt").to(model.device)
+    with torch.no_grad():
+        output = model.generate(**tokenized, max_new_tokens=300)
+    return gemma_processor.decode(output[0][tokenized["input_ids"].shape[-1]:], skip_special_tokens=True)
+
+# User preference summary
+def summarize_user_preferences(user_id, top_k=3):
+    pids = user_history.get(user_id, [])
+    rows = df[df["product_id"].isin(pids)]
+    if rows.empty:
+        return "None", "None", "None"
+    brands = rows["brand"].dropna().astype(str).value_counts().index.tolist()[:top_k]
+    styles = rows["style_attributes"].astype(str).value_counts().index.tolist()[:top_k]
+    descs = rows["meta_info"].dropna().astype(str).tolist()
+    return ", ".join(brands), ", ".join(styles), " ".join(descs[:top_k])
+
+# ========== APP STARTS ==========
+st.set_page_config("🛍️ Fashion Visual Search")
+st.title("👗 Fashion Visual Search & Outfit Assistant")
+
+image_embeddings, text_embeddings, ids, _, faiss_index, df, user_history, trend_string = load_assets()
+
+dtype = torch.float16 if torch.cuda.is_available() else torch.float32
+
+clip_model = CLIPModel.from_pretrained("openai/clip-vit-large-patch14").to(device).eval()
+clip_processor = CLIPProcessor.from_pretrained("openai/clip-vit-large-patch14", use_fast=True)
+
+text_model = SentenceTransformer('all-MiniLM-L6-v2')
+
+model_id = "google/gemma-3-4b-it"
+model = Gemma3ForConditionalGeneration.from_pretrained(model_id, torch_dtype=dtype, device_map="auto").eval()
+gemma_processor = AutoProcessor.from_pretrained(model_id)
+
+username = st.text_input("👤 Enter your username:")
+if username:
+    uploaded_image = st.file_uploader("📷 Upload a fashion image", type=["jpg", "png"])
+    text_query = st.text_input("📝 Optional: Describe what you're looking for")
+    num_results = st.slider("🔢 Number of similar items", 1, 20, 5)
+    num_suggestions = st.slider("💡 Number of outfit suggestions", 1, 10, 3)
+
+    if uploaded_image:
+        st.image(uploaded_image, caption="Uploaded Image", width=300)
+
+        img = PILImage.open(uploaded_image)
+        similar_ids = search_similar(image=img, text=text_query, top_k=num_results)
+        st.subheader("🎯 Similar Products")
+        for pid in similar_ids:
+            row = df[df["product_id"] == pid].iloc[0]
+            st.image(row["feature_image_s3"], width=200)
+            st.write(f"**{row['product_name']}** — ₹{row['selling_price']}")
+            st.write(f"Brand: {row['brand']}")
+            if username not in user_history:
+                user_history[username] = []
+            user_history[username].append(pid)
+
+        st.subheader("🧠 Outfit Suggestions")
+        top_row = df[df["product_id"] == similar_ids[0]].iloc[0]
+        suggestions = generate_outfit_gemma(img, top_row, username, suggestions=num_suggestions)
+        st.markdown(suggestions)
+
+        st.subheader("🧾 Inventory Text Search")
+        text_only_ids = search_similar(image=None, text=text_query, top_k=num_results)
+        for pid in text_only_ids:
+            row = df[df["product_id"] == pid].iloc[0]
+            st.image(row["feature_image_s3"], width=200)
+            st.write(f"{row['product_name']} — ₹{row['selling_price']}")
+            st.write(f"Brand: {row['brand']}")
 
-"""
-# Welcome to Streamlit!
-
-Edit `/streamlit_app.py` to customize this app to your heart's desire :heart:.
-If you have any questions, checkout our [documentation](https://docs.streamlit.io) and [community
-forums](https://discuss.streamlit.io).
-
-In the meantime, below is an example of what you can do with just a few lines of code:
-"""
-
-num_points = st.slider("Number of points in spiral", 1, 10000, 1100)
-num_turns = st.slider("Number of turns in spiral", 1, 300, 31)
-
-indices = np.linspace(0, 1, num_points)
-theta = 2 * np.pi * num_turns * indices
-radius = indices
-
-x = radius * np.cos(theta)
-y = radius * np.sin(theta)
-
-df = pd.DataFrame({
-    "x": x,
-    "y": y,
-    "idx": indices,
-    "rand": np.random.randn(num_points),
-})
-
-st.altair_chart(alt.Chart(df, height=700, width=700)
-    .mark_point(filled=True)
-    .encode(
-        x=alt.X("x", axis=None),
-        y=alt.Y("y", axis=None),
-        color=alt.Color("idx", legend=None, scale=alt.Scale()),
-        size=alt.Size("rand", legend=None, scale=alt.Scale(range=[1, 150])),
-    ))
+        st.subheader("📦 Personalized History-Based Suggestions")
+        brands, styles, desc = summarize_user_preferences(username)
+        if brands == "None":
+            st.warning("⚠️ No history found yet. Try uploading images first!")
+        else:
+            hist_ids = [pid for pid in ids if any(b in text_embeddings[pid] for b in brands.split(", "))][:num_results]
+            for pid in hist_ids:
+                row = df[df["product_id"] == pid].iloc[0]
+                st.image(row["feature_image_s3"], width=200)
+                st.write(f"{row['product_name']} — ₹{row['selling_price']}")
+                st.write(f"Brand: {row['brand']}")