Update app.py

app.py

@@ -58,6 +58,14 @@ def get_segmented_embedding(img, final_mask):
         image_features /= image_features.norm(dim=-1, keepdim=True)
     return image_features.cpu().numpy().flatten()
 
+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().flatten()
+
+
 def segment_clothing(img, clothes=["Hat", "Upper-clothes", "Skirt", "Pants", "Dress", "Belt", "Left-shoe", "Right-shoe", "Scarf"]):
     segments = segmenter(img)
     mask_list = []
@@ -174,8 +182,8 @@ elif st.session_state.step == 'show_results':
     st.image(original_image, caption="Original Image", use_column_width=True)
     
     # 세그먼트된 이미지에서 임베딩 추출
-    query_embedding = get_segmented_embedding(st.session_state.query_image, st.session_state.segmented_image)
-    
+    #query_embedding = get_segmented_embedding(st.session_state.query_image, st.session_state.segmented_image)
+    query_embedding = get_image_embedding(st.session_state.query_image, st.session_state.segmented_image)
     similar_images = find_similar_images(query_embedding, collection)
     
     st.subheader("Similar Items:")