Update app.py

app.py

@@ -6,10 +6,9 @@ from PIL import Image
 from io import BytesIO
 import time
 import numpy as np
-from ultralytics import YOLO
-import chromadb
 from transformers import pipeline
-from sklearn.metrics.pairwise import cosine_similarity
+import chromadb
+from sklearn.metrics.pairwise import euclidean_distances  # 유클리드 거리 계산 추가
 
 # Load segmentation model
 segmenter = pipeline(model="mattmdjaga/segformer_b2_clothes")
@@ -25,20 +24,10 @@ def load_clip_model():
 
 clip_model, preprocess_val, tokenizer, device = load_clip_model()
 
-# Load YOLOv8 model
-#@st.cache_resource
-#def load_yolo_model():
-#    return YOLO("./best.pt")
-
-#yolo_model = load_yolo_model()
-
 # Load chromaDB
 client = chromadb.PersistentClient(path="./clothesDB_11GmarketMusinsa")
-#collection = client.get_collection(name="clothes_items_ver3")
 collection = client.get_collection(name="clothes")
 
-
-
 # Helper functions
 def load_image_from_url(url, max_retries=3):
     for attempt in range(max_retries):
@@ -53,63 +42,62 @@ def load_image_from_url(url, max_retries=3):
             else:
                 return None
 
-def get_image_embedding(image):
-    image_tensor = preprocess_val(image).unsqueeze(0).to(device)
+# 세그먼트 마스크 기반 임베딩 추출
+def get_segmented_embedding(img, final_mask):
+    img_array = np.array(img)
+    final_mask_array = np.array(final_mask)
+    
+    # 마스크 적용 (배경을 흰색으로 처리)
+    img_array[final_mask_array == 0] = 255
+    masked_img = Image.fromarray(img_array)
+    
+    # 마스크된 이미지로부터 임베딩 추출
+    image_tensor = preprocess_val(masked_img).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().flatten()
 
 def segment_clothing(img, clothes=["Hat", "Upper-clothes", "Skirt", "Pants", "Dress", "Belt", "Left-shoe", "Right-shoe", "Scarf"]):
-    # Segment image
     segments = segmenter(img)
-
-    # Create list of masks
     mask_list = []
     detected_categories = []
     for s in segments:
         if s['label'] in clothes:
             mask_list.append(s['mask'])
-            detected_categories.append(s['label'])  # Store detected categories
+            detected_categories.append(s['label'])
 
-    # Paste all masks on top of each other
-    final_mask = np.zeros_like(np.array(img)[:, :, 0])  # Initialize mask
+    final_mask = np.zeros_like(np.array(img)[:, :, 0])
     for mask in mask_list:
         current_mask = np.array(mask)
-        final_mask = np.maximum(final_mask, current_mask)  # Use maximum to combine masks
+        final_mask = np.maximum(final_mask, current_mask)
 
-    # Convert final mask from np array to PIL image
-    final_mask = Image.fromarray(final_mask.astype(np.uint8) * 255)  # Convert to binary mask
-
-    # Apply mask to original image
-    img_with_alpha = img.convert("RGBA")  # Ensure the image has an alpha channel
+    final_mask = Image.fromarray(final_mask.astype(np.uint8) * 255)
+    img_with_alpha = img.convert("RGBA")
     img_with_alpha.putalpha(final_mask)
 
-    return img_with_alpha.convert("RGB"), final_mask, detected_categories  # Return detected categories
-    
+    return img_with_alpha.convert("RGB"), final_mask, detected_categories
+
+# 유클리드 거리 기반으로 유사한 이미지 찾기
 def find_similar_images(query_embedding, collection, top_k=5):
-    query_embedding = query_embedding.reshape(1, -1)  # Reshape to 2D array for ChromaDB
-    results = collection.query(
-        query_embeddings=query_embedding,
-        n_results=top_k,
-        include=['metadatas', 'distances']  
-    )
+    query_embedding = query_embedding.reshape(1, -1)
     
-    top_metadatas = results['metadatas'][0]  
-    top_distances = results['distances'][0]  
+    all_embeddings = np.array(collection.get(include=['embeddings'])['embeddings'])
+    distances = euclidean_distances(query_embedding, all_embeddings).flatten()  # 유클리드 거리 계산
+    
+    top_indices = np.argsort(distances)[:top_k]
+    top_metadatas = [collection.get(ids=[str(idx)])['metadatas'][0] for idx in top_indices]
+    top_distances = distances[top_indices]
     
     structured_results = []
     for metadata, distance in zip(top_metadatas, top_distances):
         structured_results.append({
             'info': metadata,
-            'similarity': 1 - distance  
+            'similarity': 1 / (1 + distance)  # 유클리드 거리가 작을수록 유사도가 높음
         })
     
     return structured_results
 
-
-
-
 # 세션 상태 초기화
 if 'step' not in st.session_state:
     st.session_state.step = 'input'
@@ -117,7 +105,7 @@ 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 'segmented_image' not in st.session_state:  # Add segmented_image to session state
+if 'segmented_image' not in st.session_state:
     st.session_state.segmented_image = None
 if 'selected_category' not in st.session_state:
     st.session_state.selected_category = None
@@ -125,7 +113,6 @@ if 'selected_category' not in st.session_state:
 # 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"):
@@ -133,10 +120,9 @@ if st.session_state.step == 'input':
             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
-                # Perform segmentation
                 segmented_image, final_mask, detected_categories = segment_clothing(query_image)
-                st.session_state.segmented_image = segmented_image  # Store segmented image in session state
-                st.session_state.detections = detected_categories  # Store detected categories
+                st.session_state.segmented_image = segmented_image
+                st.session_state.detections = detected_categories
                 st.image(segmented_image, caption="Segmented Image", use_column_width=True)
                 if st.session_state.detections:
                     st.session_state.step = 'select_category'
@@ -160,19 +146,18 @@ elif st.session_state.step == 'select_category':
         st.warning("No categories detected.")
 
 elif st.session_state.step == 'show_results':
-    original_image = st.session_state.query_image.convert("RGB")  # Convert to RGB before displaying
+    original_image = st.session_state.query_image.convert("RGB")
     st.image(original_image, caption="Original Image", use_column_width=True)
     
-    # Get the embedding of the segmented image
-    query_embedding = get_image_embedding(st.session_state.segmented_image)  # Use the segmented image from session state
-
+    # 세그먼트된 이미지에서 임베딩 추출
+    query_embedding = get_segmented_embedding(st.session_state.query_image, st.session_state.segmented_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['image_url'], use_column_width=True)
             st.image(img['info']['image_url'], use_column_width=True)
         with col2:
             st.write(f"Name: {img['info']['name']}")
@@ -188,4 +173,4 @@ elif st.session_state.step == 'show_results':
         st.session_state.step = 'input'
         st.session_state.query_image_url = ''
         st.session_state.detections = []
-        st.session_state.segmented_image = None  # Reset segmented_image
+        st.session_state.segmented_image = None