Update app.py

app.py

@@ -52,12 +52,13 @@ def load_image_from_url(url, max_retries=3):
             else:
                 return None
 
+# 기존의 get_image_embedding 함수도 수정
 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()
+    return image_features.cpu().numpy().squeeze().tolist()  # numpy 배열을 파이썬 리스트로 변환
 
 def get_text_embedding(text):
     text_tokens = tokenizer([text]).to(device)
@@ -65,7 +66,6 @@ def get_text_embedding(text):
         text_features = clip_model.encode_text(text_tokens)
         text_features /= text_features.norm(dim=-1, keepdim=True)
     return text_features.cpu().numpy()
-
 def get_average_embedding(main_image_url, additional_image_urls):
     embeddings = []
     
@@ -81,10 +81,52 @@ def get_average_embedding(main_image_url, additional_image_urls):
             embeddings.append(get_image_embedding(img))
     
     if embeddings:
-        return np.mean(embeddings, axis=0)
+        avg_embedding = np.mean(embeddings, axis=0)
+        return avg_embedding.squeeze().tolist()  # numpy 배열을 파이썬 리스트로 변환
     else:
         return None
 
+def update_collection_embeddings():
+    all_ids = collection.get(include=['metadatas'])['ids']
+    all_metadata = collection.get(include=['metadatas'])['metadatas']
+    
+    batch_size = 100  # 한 번에 처리할 아이템 수
+    for i in range(0, len(all_ids), batch_size):
+        batch_ids = all_ids[i:i+batch_size]
+        batch_metadata = all_metadata[i:i+batch_size]
+        
+        batch_embeddings = []
+        valid_ids = []
+        
+        for id, metadata in zip(batch_ids, batch_metadata):
+            main_image_url = metadata['image_url']
+            additional_image_urls = metadata.get('additional_images', [])
+            
+            try:
+                avg_embedding = get_average_embedding(main_image_url, additional_image_urls)
+                if avg_embedding is not None:
+                    batch_embeddings.append(avg_embedding)  # 이미 리스트 형태로 반환됨
+                    valid_ids.append(id)
+                else:
+                    st.warning(f"Failed to generate embedding for item {id}")
+            except Exception as e:
+                st.error(f"Error processing item {id}: {str(e)}")
+        
+        if valid_ids:
+            try:
+                collection.update(
+                    ids=valid_ids,
+                    embeddings=batch_embeddings
+                )
+                st.success(f"Updated embeddings for {len(valid_ids)} items")
+            except Exception as e:
+                st.error(f"Error updating embeddings: {str(e)}")
+                st.error(f"First embedding type: {type(batch_embeddings[0])}")
+                st.error(f"First embedding shape: {len(batch_embeddings[0])}")
+        
+        # 진행 상황 표시
+        st.progress((i + batch_size) / len(all_ids))
+
 def find_similar_images(query_embedding, collection, top_k=5):
     results = collection.query(
         query_embeddings=[query_embedding.squeeze().tolist()],