Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,40 +1,117 @@
-import altair as alt
-import numpy as np
-import pandas as pd
 import streamlit as st
+import boto3
+from botocore.exceptions import NoCredentialsError
+from io import BytesIO
+from PIL import Image
+import pandas as pd
+import matplotlib.pyplot as plt
+from PIL import Image
+import os
+import faiss
+import pickle
+import torch
+from transformers import CLIPModel, CLIPProcessor
+from huggingface_hub import hf_hub_download
+import json
+
+# Lấy thông tin từ Streamlit Secrets
+# aws_access_key = st.secrets["AWS_ACCESS_KEY_ID"]
+# aws_secret_key = st.secrets["AWS_SECRET_ACCESS_KEY"]
+# aws_access_key = 'AKIATS5GX2D62YHYRFWL'
+# aws_secret_key = '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'
+
+
+# Khởi tạo client S3 với thông tin cấu hình từ secrets
+s3 = boto3.client(
+    's3'
+)
+
+def get_image_from_s3(bucket_name, img_id):
+    # object_name = str(img_id) + '.jpg'
+
+    try:
+        # # Tải ảnh từ S3
+        # response = s3.get_object(Bucket=bucket_name, Key=object_name)
+        # img_data = response['Body'].read()  # Đọc dữ liệu ảnh
+
+        # # Chuyển đổi dữ liệu ảnh thành ảnh PIL
+        # img = Image.open(BytesIO(img_data))
+        # return img
+        
+        return f"https://{bucket_name}.s3.amazonaws.com/{img_id}.jpg"
+        
+    except NoCredentialsError:
+        st.error("Credentials not available.")
+        return None
+    except Exception as e:
+        st.error(f"Error fetching image: {e}")
+        return None
+
+def show_img(img_id, score):
+    # Lấy ảnh từ S3
+    img = get_image_from_s3(bucket_name, img_id)
+
+    if img:
+        img_style = style[style['id'] == img_id]
+        
+        if not img_style.empty:
+            parts = []
+            parts.append(str(img_style['gender'].values[0]))
+            parts.append(str(img_style['masterCategory'].values[0]))
+            parts.append(str(img_style['subCategory'].values[0]))
+            parts.append(str(img_style['articleType'].values[0]))
+            parts.append(str(img_style['baseColour'].values[0]))
+            parts.append(str(img_style['year'].values[0]))
+            parts.append(str(img_style['usage'].values[0]))
+            parts.append(str(img_style['productDisplayName'].values[0]))
+
+            text = '- '.join(parts)
+            text += f'\n Score: {score}'
+            st.image(img, caption=text, use_container_width=True)
+
+
+
+def search_faiss(model, processor, index, id_map, prompt, top_k=5, device='cpu'):
+    inputs = processor(text=[prompt], return_tensors='pt', padding=True).to(device)
+    with torch.no_grad():
+        txt_emb = model.get_text_features(input_ids=inputs['input_ids'], attention_mask=inputs['attention_mask'])
+        txt_emb = txt_emb / txt_emb.norm(p=2, dim=-1, keepdim=True)
+    q = txt_emb.cpu().numpy().astype('float32')
+
+    D, I = index.search(q, top_k)
+    return [(id_map[i], float(D[0][j])) for j, i in enumerate(I[0])]
+
+
+def running(prompt, top_k=5):
+    results = search_faiss(
+        model, processor,
+        index, id_map,
+        prompt=prompt,
+        top_k=top_k,
+    )
+
+    for img_id, score in results:
+        show_img(img_id, score)
+   
+   
+style = pd.read_csv('styles.csv', usecols=range(10))
+bucket_name = "image-text-retrieval"  # Tên bucket của bạn
+your_username = 'roy214'      
+
+ 
+# Load model từ Hugging Face Hub
+model = CLIPModel.from_pretrained(f"{your_username}/clip-finetuned-fashion").to("cpu").eval()
+processor = CLIPProcessor.from_pretrained(f"{your_username}/clip-finetuned-fashion")
+
+# Load FAISS index và id_map
+index_path = hf_hub_download(repo_id=f"{your_username}/clip-finetuned-fashion", filename="faiss_index.bin")
+mapping_path = hf_hub_download(repo_id=f"{your_username}/clip-finetuned-fashion", filename="id_map.json")
+
+index = faiss.read_index(index_path)
+
+with open(mapping_path, "r") as f:
+    id_map = json.load(f)
 
-"""
-# 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])),
-    ))
+# show_img(59403, 19)
+st.text("Enter prompt")
+running("Dress Women Apparel Red", top_k=5)