Update app.py

app.py

@@ -1,209 +1,224 @@
-import streamlit as st
-import open_clip
-import torch
-import requests
-from PIL import Image
-from io import BytesIO
-import time
-import json
-import numpy as np
-import onnxruntime as ort
-import cv2
-import chromadb
-
-@st.cache_resource
-def load_clip_model():
-    model, preprocess_train, preprocess_val = open_clip.create_model_and_transforms('hf-hub:Marqo/marqo-fashionSigLIP')
-    tokenizer = open_clip.get_tokenizer('hf-hub:Marqo/marqo-fashionSigLIP')
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    model.to(device)
-    return model, preprocess_val, tokenizer, device
-
-clip_model, preprocess_val, tokenizer, device = load_clip_model()
-
-@st.cache_resource
-def load_onnx_model():
-    session = ort.InferenceSession("./accessary_weights.onnx")
-    return session
-
-onnx_session = load_onnx_model()
-
-def load_image_from_url(url, max_retries=3):
-    for attempt in range(max_retries):
-        try:
-            response = requests.get(url, timeout=10)
-            response.raise_for_status()
-            img = Image.open(BytesIO(response.content)).convert('RGB')
-            return img
-        except (requests.RequestException, Image.UnidentifiedImageError) as e:
-            if attempt < max_retries - 1:
-                time.sleep(1)
-            else:
-                return None
-
-client = chromadb.PersistentClient(path="./accessaryDB")
-collection = client.get_collection(name="accessary_items_ver2")
-
-def get_image_embedding(image):
-    image_tensor = preprocess_val(image).unsqueeze(0).to(device)
-    with torch.no_grad():
-        image_features = clip_model.encode_image(image_tensor)
-        image_features /= image_features.norm(dim=-1, keepdim=True)
-    return image_features.cpu().numpy()
-
-def get_text_embedding(text):
-    text_tokens = tokenizer([text]).to(device)
-    with torch.no_grad():
-        text_features = clip_model.encode_text(text_tokens)
-        text_features /= text_features.norm(dim=-1, keepdim=True)
-    return text_features.cpu().numpy()
-
-def get_all_embeddings_from_collection(collection):
-    all_embeddings = collection.get(include=['embeddings'])['embeddings']
-    return np.array(all_embeddings)
-
-def get_metadata_from_ids(collection, ids):
-    results = collection.get(ids=ids)
-    return results['metadatas']
-
-def find_similar_images(query_embedding, collection, top_k=5):
-    database_embeddings = get_all_embeddings_from_collection(collection)
-    similarities = np.dot(database_embeddings, query_embedding.T).squeeze()
-    top_indices = np.argsort(similarities)[::-1][:top_k]
-    
-    all_data = collection.get(include=['metadatas'])['metadatas']
-    
-    top_metadatas = [all_data[idx] for idx in top_indices]
-    
-    results = []
-    for idx, metadata in enumerate(top_metadatas):
-        results.append({
-            'info': metadata,
-            'similarity': similarities[top_indices[idx]]
-        })
-    return results
-
-def detect_clothing_onnx(image):
-    input_image = np.array(image.resize((640, 640)), dtype=np.float32)
-    input_image = np.transpose(input_image, [2, 0, 1]) 
-    input_image = np.expand_dims(input_image, axis=0)  
-    input_image /= 255.0  
-
-    inputs = {onnx_session.get_inputs()[0].name: input_image}
-    outputs = onnx_session.run(None, inputs)
-    
-    detections = outputs[0]  
-    categories = []
-    for detection in detections:
-        x1, y1, x2, y2, conf, cls = detection
-        category = str(int(cls))  
-        if category in ['Bracelets', 'Broches', 'belt', 'earring', 'maangtika', 'necklace', 'nose ring', 'ring', 'tiara']:
-            categories.append({
-                'category': category,
-                'bbox': [int(x1), int(y1), int(x2), int(y2)],
-                'confidence': conf
-            })
-    return categories
-
-def crop_image(image, bbox):
-    return image.crop((bbox[0], bbox[1], bbox[2], bbox[3]))
-
-# 세션 상태 초기화
-if 'step' not in st.session_state:
-    st.session_state.step = 'input'
-if 'query_image_url' not in st.session_state:
-    st.session_state.query_image_url = ''
-if 'detections' not in st.session_state:
-    st.session_state.detections = []
-if 'selected_category' not in st.session_state:
-    st.session_state.selected_category = None
-
-# Streamlit app
-st.title("Advanced Fashion Search App")
-
-# 단계별 처리
-if st.session_state.step == 'input':
-    st.session_state.query_image_url = st.text_input("Enter image URL:", st.session_state.query_image_url)
-    if st.button("Detect Clothing"):
-        if st.session_state.query_image_url:
-            query_image = load_image_from_url(st.session_state.query_image_url)
-            if query_image is not None:
-                st.session_state.query_image = query_image
-                st.session_state.detections = detect_clothing_onnx(query_image)
-                if st.session_state.detections:
-                    st.session_state.step = 'select_category'
-                else:
-                    st.warning("No clothing items detected in the image.")
-            else:
-                st.error("Failed to load the image. Please try another URL.")
-        else:
-            st.warning("Please enter an image URL.")
-
-# Update the 'select_category' step
-elif st.session_state.step == 'select_category':
-    st.image(st.session_state.query_image, caption="Query Image", use_column_width=True)
-    st.subheader("Detected Clothing Items:")
-    
-    for detection in st.session_state.detections:
-        col1, col2 = st.columns([1, 3])
-        with col1:
-            st.write(f"{detection['category']} (Confidence: {detection['confidence']:.2f})")
-        with col2:
-            cropped_image = crop_image(st.session_state.query_image, detection['bbox'])
-            st.image(cropped_image, caption=detection['category'], use_column_width=True)
-    
-    options = [f"{d['category']} (Confidence: {d['confidence']:.2f})" for d in st.session_state.detections]
-    selected_option = st.selectbox("Select a category to search:", options)
-    
-    if st.button("Search Similar Items"):
-        st.session_state.selected_category = selected_option
-        st.session_state.step = 'show_results'
-
-elif st.session_state.step == 'show_results':
-    st.image(st.session_state.query_image, caption="Query Image", use_column_width=True)
-    selected_detection = next(d for d in st.session_state.detections 
-                              if f"{d['category']} (Confidence: {d['confidence']:.2f})" == st.session_state.selected_category)
-    cropped_image = crop_image(st.session_state.query_image, selected_detection['bbox'])
-    st.image(cropped_image, caption="Cropped Image", use_column_width=True)
-    query_embedding = get_image_embedding(cropped_image)
-    similar_images = find_similar_images(query_embedding, collection)
-    
-    st.subheader("Similar Items:")
-    for img in similar_images:
-        col1, col2 = st.columns(2)
-        with col1:
-            st.image(img['info']['image_url'], use_column_width=True)
-        with col2:
-            st.write(f"Name: {img['info']['name']}")
-            st.write(f"Brand: {img['info']['brand']}")
-            st.write(f"Category: {img['info']['category']}")
-            st.write(f"Price: {img['info']['price']}")
-            st.write(f"Discount: {img['info']['discount']}%")
-            st.write(f"Similarity: {img['similarity']:.2f}")
-    
-    if st.button("Start New Search"):
-        st.session_state.step = 'input'
-        st.session_state.query_image_url = ''
-        st.session_state.detections = []
-        st.session_state.selected_category = None
-
-else:  # Text search
-    query_text = st.text_input("Enter search text:")
-    if st.button("Search by Text"):
-        if query_text:
-            text_embedding = get_text_embedding(query_text)
-            similar_images = find_similar_images(text_embedding, collection)
-            st.subheader("Similar Items:")
-            for img in similar_images:
-                col1, col2 = st.columns(2)
-                with col1:
-                    st.image(img['info']['image_url'], use_column_width=True)
-                with col2:
-                    st.write(f"Name: {img['info']['name']}")
-                    st.write(f"Brand: {img['info']['brand']}")
-                    st.write(f"Category: {img['info']['category']}")
-                    st.write(f"Price: {img['info']['price']}")
-                    st.write(f"Discount: {img['info']['discount']}%")
-                    st.write(f"Similarity: {img['similarity']:.2f}")
-        else:
+import streamlit as st
+import open_clip
+import torch
+import requests
+from PIL import Image
+from io import BytesIO
+import time
+import json
+import numpy as np
+import onnxruntime as ort
+import cv2
+import chromadb
+
+@st.cache_resource
+def load_clip_model():
+    model, preprocess_train, preprocess_val = open_clip.create_model_and_transforms('hf-hub:Marqo/marqo-fashionSigLIP')
+    tokenizer = open_clip.get_tokenizer('hf-hub:Marqo/marqo-fashionSigLIP')
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model.to(device)
+    return model, preprocess_val, tokenizer, device
+
+clip_model, preprocess_val, tokenizer, device = load_clip_model()
+
+@st.cache_resource
+def load_onnx_model():
+    session = ort.InferenceSession("./accessary_weights.onnx")
+    return session
+
+onnx_session = load_onnx_model()
+
+def load_image_from_url(url, max_retries=3):
+    for attempt in range(max_retries):
+        try:
+            response = requests.get(url, timeout=10)
+            response.raise_for_status()
+            img = Image.open(BytesIO(response.content)).convert('RGB')
+            return img
+        except (requests.RequestException, Image.UnidentifiedImageError) as e:
+            if attempt < max_retries - 1:
+                time.sleep(1)
+            else:
+                return None
+
+client = chromadb.PersistentClient(path="./accessaryDB")
+collection = client.get_collection(name="accessary_items_ver2")
+
+def get_image_embedding(image):
+    image_tensor = preprocess_val(image).unsqueeze(0).to(device)
+    with torch.no_grad():
+        image_features = clip_model.encode_image(image_tensor)
+        image_features /= image_features.norm(dim=-1, keepdim=True)
+    return image_features.cpu().numpy()
+
+def get_text_embedding(text):
+    text_tokens = tokenizer([text]).to(device)
+    with torch.no_grad():
+        text_features = clip_model.encode_text(text_tokens)
+        text_features /= text_features.norm(dim=-1, keepdim=True)
+    return text_features.cpu().numpy()
+
+def get_all_embeddings_from_collection(collection):
+    all_embeddings = collection.get(include=['embeddings'])['embeddings']
+    return np.array(all_embeddings)
+
+def get_metadata_from_ids(collection, ids):
+    results = collection.get(ids=ids)
+    return results['metadatas']
+
+def find_similar_images(query_embedding, collection, top_k=5):
+    database_embeddings = get_all_embeddings_from_collection(collection)
+    similarities = np.dot(database_embeddings, query_embedding.T).squeeze()
+    top_indices = np.argsort(similarities)[::-1][:top_k]
+    
+    all_data = collection.get(include=['metadatas'])['metadatas']
+    
+    top_metadatas = [all_data[idx] for idx in top_indices]
+    
+    results = []
+    for idx, metadata in enumerate(top_metadatas):
+        results.append({
+            'info': metadata,
+            'similarity': similarities[top_indices[idx]]
+        })
+    return results
+
+onnx_model_labels = ['Bracelets', 'Broches', 'belt', 'earring', 'maangtika', 'necklace', 'nose ring', 'ring', 'tiara']
+
+def preprocess_for_onnx(image, input_size=(640, 640)):
+    # 이미지 크기 조정 (ONNX 모델의 입력 크기에 맞춰 리사이즈)
+    resized_image = image.resize(input_size)
+    
+    # 이미지를 NumPy 배열로 변환하고, 0~1 사이의 값으로 정규화 (필요한 경우)
+    image_np = np.array(resized_image).astype(np.float32) / 255.0
+    
+    # 모델이 기대하는 순서대로 차원 변경 (예: HWC -> CHW)
+    image_np = np.transpose(image_np, (2, 0, 1))  # 채널 순서를 변경 (HWC -> CHW)
+    
+    # 배치 차원을 추가 (ONNX 모델은 보통 [batch, channel, height, width] 형식을 요구)
+    input_tensor = np.expand_dims(image_np, axis=0)
+    
+    return input_tensor
+
+def detect_clothing_onnx(image):
+    # ONNX 모델로 이미지에서 객체 탐지 수행
+    input_tensor = preprocess_for_onnx(image)
+    outputs = onnx_session.run(None, {onnx_session.get_inputs()[0].name: input_tensor})
+    
+    # 탐지된 객체에 대한 좌표 및 클래스 정보 추출
+    detections = outputs[0]  # 모델의 출력 형식에 맞게 수정
+    
+    categories = []
+    for detection in detections:
+        x1, y1, x2, y2, conf, cls = detection
+        category = onnx_model_labels[int(cls)]  # 클래스 인덱스를 카테고리 이름으로 변환
+        if category in ['Bracelets', 'Broches', 'belt', 'earring', 'maangtika', 'necklace', 'nose ring', 'ring', 'tiara'] : 
+                categories.append({
+                'category': category,
+                'bbox': [int(x1), int(y1), int(x2), int(y2)],
+                'confidence': conf
+            })
+    return categories
+
+def crop_image(image, bbox):
+    return image.crop((bbox[0], bbox[1], bbox[2], bbox[3]))
+
+# 세션 상태 초기화
+if 'step' not in st.session_state:
+    st.session_state.step = 'input'
+if 'query_image_url' not in st.session_state:
+    st.session_state.query_image_url = ''
+if 'detections' not in st.session_state:
+    st.session_state.detections = []
+if 'selected_category' not in st.session_state:
+    st.session_state.selected_category = None
+
+# Streamlit app
+st.title("Advanced Fashion Search App")
+
+# 단계별 처리
+if st.session_state.step == 'input':
+    st.session_state.query_image_url = st.text_input("Enter image URL:", st.session_state.query_image_url)
+    if st.button("Detect Clothing"):
+        if st.session_state.query_image_url:
+            query_image = load_image_from_url(st.session_state.query_image_url)
+            if query_image is not None:
+                st.session_state.query_image = query_image
+                st.session_state.detections = detect_clothing_onnx(query_image)
+                if st.session_state.detections:
+                    st.session_state.step = 'select_category'
+                else:
+                    st.warning("No clothing items detected in the image.")
+            else:
+                st.error("Failed to load the image. Please try another URL.")
+        else:
+            st.warning("Please enter an image URL.")
+
+# Update the 'select_category' step
+elif st.session_state.step == 'select_category':
+    st.image(st.session_state.query_image, caption="Query Image", use_column_width=True)
+    st.subheader("Detected Clothing Items:")
+    
+    for detection in st.session_state.detections:
+        col1, col2 = st.columns([1, 3])
+        with col1:
+            st.write(f"{detection['category']} (Confidence: {detection['confidence']:.2f})")
+        with col2:
+            cropped_image = crop_image(st.session_state.query_image, detection['bbox'])
+            st.image(cropped_image, caption=detection['category'], use_column_width=True)
+    
+    options = [f"{d['category']} (Confidence: {d['confidence']:.2f})" for d in st.session_state.detections]
+    selected_option = st.selectbox("Select a category to search:", options)
+    
+    if st.button("Search Similar Items"):
+        st.session_state.selected_category = selected_option
+        st.session_state.step = 'show_results'
+
+elif st.session_state.step == 'show_results':
+    st.image(st.session_state.query_image, caption="Query Image", use_column_width=True)
+    selected_detection = next(d for d in st.session_state.detections 
+                              if f"{d['category']} (Confidence: {d['confidence']:.2f})" == st.session_state.selected_category)
+    cropped_image = crop_image(st.session_state.query_image, selected_detection['bbox'])
+    st.image(cropped_image, caption="Cropped Image", use_column_width=True)
+    query_embedding = get_image_embedding(cropped_image)
+    similar_images = find_similar_images(query_embedding, collection)
+    
+    st.subheader("Similar Items:")
+    for img in similar_images:
+        col1, col2 = st.columns(2)
+        with col1:
+            st.image(img['info']['image_url'], use_column_width=True)
+        with col2:
+            st.write(f"Name: {img['info']['name']}")
+            st.write(f"Brand: {img['info']['brand']}")
+            st.write(f"Category: {img['info']['category']}")
+            st.write(f"Price: {img['info']['price']}")
+            st.write(f"Discount: {img['info']['discount']}%")
+            st.write(f"Similarity: {img['similarity']:.2f}")
+    
+    if st.button("Start New Search"):
+        st.session_state.step = 'input'
+        st.session_state.query_image_url = ''
+        st.session_state.detections = []
+        st.session_state.selected_category = None
+
+else:  # Text search
+    query_text = st.text_input("Enter search text:")
+    if st.button("Search by Text"):
+        if query_text:
+            text_embedding = get_text_embedding(query_text)
+            similar_images = find_similar_images(text_embedding, collection)
+            st.subheader("Similar Items:")
+            for img in similar_images:
+                col1, col2 = st.columns(2)
+                with col1:
+                    st.image(img['info']['image_url'], use_column_width=True)
+                with col2:
+                    st.write(f"Name: {img['info']['name']}")
+                    st.write(f"Brand: {img['info']['brand']}")
+                    st.write(f"Category: {img['info']['category']}")
+                    st.write(f"Price: {img['info']['price']}")
+                    st.write(f"Discount: {img['info']['discount']}%")
+                    st.write(f"Similarity: {img['similarity']:.2f}")
+        else:
             st.warning("Please enter a search text.")