Update app.py

app.py

@@ -11,6 +11,7 @@ import onnxruntime as ort
 import cv2
 import chromadb
 
+# CLIP 모델 로드
 @st.cache_resource
 def load_clip_model():
     model, preprocess_train, preprocess_val = open_clip.create_model_and_transforms('hf-hub:Marqo/marqo-fashionSigLIP')
@@ -21,6 +22,7 @@ def load_clip_model():
 
 clip_model, preprocess_val, tokenizer, device = load_clip_model()
 
+# ONNX 모델 로드
 @st.cache_resource
 def load_onnx_model():
     session = ort.InferenceSession("./accessary_weights.onnx")
@@ -28,6 +30,7 @@ def load_onnx_model():
 
 onnx_session = load_onnx_model()
 
+# URL에서 이미지 로드
 def load_image_from_url(url, max_retries=3):
     for attempt in range(max_retries):
         try:
@@ -41,9 +44,11 @@ def load_image_from_url(url, max_retries=3):
             else:
                 return None
 
+# ChromaDB 클라이언트 설정
 client = chromadb.PersistentClient(path="./accessaryDB")
 collection = client.get_collection(name="accessary_items_ver2")
 
+# CLIP 이미지 임베딩 추출
 def get_image_embedding(image):
     image_tensor = preprocess_val(image).unsqueeze(0).to(device)
     with torch.no_grad():
@@ -51,6 +56,7 @@ def get_image_embedding(image):
         image_features /= image_features.norm(dim=-1, keepdim=True)
     return image_features.cpu().numpy()
 
+# CLIP 텍스트 임베딩 추출
 def get_text_embedding(text):
     text_tokens = tokenizer([text]).to(device)
     with torch.no_grad():
@@ -58,14 +64,17 @@ def get_text_embedding(text):
         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)
 
+# ID를 통해 메타데이터 가져오기
 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()
@@ -85,41 +94,35 @@ def find_similar_images(query_embedding, collection, top_k=5):
 
 onnx_model_labels = ['Bracelets', 'Broches', 'belt', 'earring', 'maangtika', 'necklace', 'nose ring', 'ring', 'tiara']
 
+# ONNX 모델에 맞춘 전처리 함수
 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] 형식을 요구)
+    image_np = np.transpose(image_np, (2, 0, 1))
     input_tensor = np.expand_dims(image_np, axis=0)
-    
     return input_tensor
 
+# 의류 탐지
 def detect_clothing_onnx(image):
-    # ONNX 모델로 이미지에서 객체 탐지 수행
-    input_tensor = preprocess_for_onnx(image)
+    input_tensor = preprocess_for_onnx(image)  # 전처리 함수 호출
     outputs = onnx_session.run(None, {onnx_session.get_inputs()[0].name: input_tensor})
     
-    # 탐지된 객체에 대한 좌표 및 클래스 정보 추출
-    detections = outputs[0]  # 모델의 출력 형식에 맞게 수정
-    
+    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({
+        x1, y1, x2, y2, conf, cls = detection[:6]  # 필요한 값만 추출
+        if conf > 0.5:  # 신뢰도 임계값 설정
+            category = onnx_model_labels[int(cls)]
+            categories.append({
                 'category': category,
-                'bbox': [int(x1), int(y1), int(x2), int(y2)],
+                'bbox': [x1, y1, x2, y2],
                 'confidence': conf
             })
+    
     return categories
 
+# 이미지 자르기
 def crop_image(image, bbox):
     return image.crop((bbox[0], bbox[1], bbox[2], bbox[3]))
 
@@ -154,7 +157,6 @@ if st.session_state.step == 'input':
         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:")
@@ -202,6 +204,7 @@ elif st.session_state.step == 'show_results':
         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"):