import streamlit as st import open_clip import torch from PIL import Image import numpy as np from transformers import AutoImageProcessor, AutoModelForSemanticSegmentation import chromadb import logging import io import requests from concurrent.futures import ThreadPoolExecutor from chromadb.utils.embedding_functions import OpenCLIPEmbeddingFunction from chromadb.utils.data_loaders import ImageLoader # 로깅 설정 logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) class CustomFashionEmbeddingFunction: def __init__(self): self.model, _, self.preprocess = open_clip.create_model_and_transforms('hf-hub:Marqo/marqo-fashionSigLIP') self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") self.model = self.model.to(self.device) def __call__(self, input): try: # 입력이 URL이나 경로인 경우 처리 processed_images = [] for img in input: if isinstance(img, (str, bytes)): if isinstance(img, str): response = requests.get(img) img = Image.open(io.BytesIO(response.content)).convert('RGB') else: img = Image.open(io.BytesIO(img)).convert('RGB') elif isinstance(img, np.ndarray): img = Image.fromarray(img.astype('uint8')).convert('RGB') processed_img = self.preprocess(img).unsqueeze(0).to(self.device) processed_images.append(processed_img) # 배치 처리 batch = torch.cat(processed_images) # CLIP 임베딩 추출 with torch.no_grad(): clip_features = self.model.encode_image(batch) clip_features = clip_features.cpu().numpy() # 색상 특징 추출 color_features_list = [] for img in input: if isinstance(img, (str, bytes)): if isinstance(img, str): response = requests.get(img) img = Image.open(io.BytesIO(response.content)).convert('RGB') else: img = Image.open(io.BytesIO(img)).convert('RGB') elif isinstance(img, np.ndarray): img = Image.fromarray(img.astype('uint8')).convert('RGB') color_features = self.extract_color_histogram(img) color_features_list.append(color_features) # 특징 결합 combined_embeddings = [] for clip_emb, color_feat in zip(clip_features, color_features_list): # CLIP 임베딩을 768차원으로 패딩 if clip_emb.shape[0] < 768: padding = np.zeros(768 - clip_emb.shape[0]) clip_emb = np.concatenate([clip_emb, padding]) else: clip_emb = clip_emb[:768] # 768차원으로 자르기 # 색상 특징을 768차원으로 확장 color_features_expanded = np.repeat(color_feat, 32) # 24 * 32 = 768 # 정규화 clip_emb = clip_emb / (np.linalg.norm(clip_emb) + 1e-8) color_features_expanded = color_features_expanded / (np.linalg.norm(color_features_expanded) + 1e-8) # 가중치 결합 combined = clip_emb * 0.7 + color_features_expanded * 0.3 combined = combined / (np.linalg.norm(combined) + 1e-8) combined_embeddings.append(combined) return np.array(combined_embeddings) except Exception as e: logger.error(f"Error in embedding function: {e}") raise def extract_color_histogram(self, image): """Extract color histogram from the image""" try: if isinstance(image, (str, bytes)): if isinstance(image, str): response = requests.get(image) image = Image.open(io.BytesIO(response.content)) else: image = Image.open(io.BytesIO(image)) if not isinstance(image, np.ndarray): img_array = np.array(image) else: img_array = image # HSV 변환 및 히스토그램 계산 img_hsv = Image.fromarray(img_array.astype('uint8')).convert('HSV') hsv_pixels = np.array(img_hsv) h_hist = np.histogram(hsv_pixels[:,:,0], bins=8, range=(0, 256))[0] s_hist = np.histogram(hsv_pixels[:,:,1], bins=8, range=(0, 256))[0] v_hist = np.histogram(hsv_pixels[:,:,2], bins=8, range=(0, 256))[0] # 정규화 h_hist = h_hist / (h_hist.sum() + 1e-8) s_hist = s_hist / (s_hist.sum() + 1e-8) v_hist = v_hist / (v_hist.sum() + 1e-8) return np.concatenate([h_hist, s_hist, v_hist]) except Exception as e: logger.error(f"Color histogram extraction error: {e}") return np.zeros(24) # Initialize session state if 'image' not in st.session_state: st.session_state.image = None if 'detected_items' not in st.session_state: st.session_state.detected_items = None if 'selected_item_index' not in st.session_state: st.session_state.selected_item_index = None if 'upload_state' not in st.session_state: st.session_state.upload_state = 'initial' if 'search_clicked' not in st.session_state: st.session_state.search_clicked = False # Load segmentation model @st.cache_resource def load_segmentation_model(): try: model_name = "mattmdjaga/segformer_b2_clothes" image_processor = AutoImageProcessor.from_pretrained(model_name) model = AutoModelForSemanticSegmentation.from_pretrained(model_name) if torch.cuda.is_available(): model = model.to('cuda') return model, image_processor except Exception as e: logger.error(f"Error loading segmentation model: {e}") raise # ChromaDB 설정 def setup_multimodal_collection(): """멀티모달 컬렉션 설정""" try: client = chromadb.PersistentClient(path="./fashion_multimodal_db_original") embedding_function = CustomFashionEmbeddingFunction() data_loader = ImageLoader() # 기존 컬렉션 가져오기 try: collection = client.get_collection( name="fashion_multimodal", embedding_function=embedding_function, data_loader=data_loader ) logger.info("Successfully connected to existing clothes_multimodal collection") return collection except Exception as e: logger.error(f"Error getting existing collection: {e}") # 컬렉션이 없는 경우에만 새로 생성 collection = client.create_collection( name="clothes_multimodal", embedding_function=embedding_function, data_loader=data_loader ) logger.info("Created new clothes_multimodal collection") return collection except Exception as e: logger.error(f"Error setting up multimodal collection: {e}") raise def process_segmentation(image): """Segmentation processing""" try: model, image_processor = load_segmentation_model() # 이미지 전처리 inputs = image_processor(image, return_tensors="pt") if torch.cuda.is_available(): inputs = {k: v.to('cuda') for k, v in inputs.items()} # 추론 with torch.no_grad(): outputs = model(**inputs) # 로직 및 후처리 logits = outputs.logits.cpu() upsampled_logits = torch.nn.functional.interpolate( logits, size=image.size[::-1], # (height, width) mode="bilinear", align_corners=False, ) # 세그멘테이션 마스크 생성 seg_masks = upsampled_logits.argmax(dim=1).numpy() processed_items = [] unique_labels = np.unique(seg_masks) for label_idx in unique_labels: if label_idx == 0: # background continue mask = (seg_masks[0] == label_idx).astype(float) processed_segment = { 'label': f"Item_{label_idx}", # 라벨 매핑이 필요하다면 여기서 처리 'score': 1.0, # confidence score 계산이 필요하다면 추가 'mask': mask } processed_items.append(processed_segment) logger.info(f"Successfully processed {len(processed_items)} segments") return processed_items except Exception as e: logger.error(f"Segmentation error: {str(e)}") import traceback logger.error(traceback.format_exc()) return [] def search_similar_items(image, mask=None, top_k=10): """멀티모달 검색 수행""" try: collection = setup_multimodal_collection() # 마스크 적용 if mask is not None: mask_3d = np.stack([mask] * 3, axis=-1) masked_image = np.array(image) * mask_3d query_image = Image.fromarray(masked_image.astype(np.uint8)) else: query_image = image # 검색 수행 results = collection.query( query_images=[np.array(query_image)], n_results=top_k, include=['metadatas', 'distances'] ) if not results or 'metadatas' not in results: return [] similar_items = [] for metadata, distance in zip(results['metadatas'][0], results['distances'][0]): # L2 거리를 코사인 유사도로 변환 # 정규화된 벡터 간의 L2 거리(d)와 코사인 유사도(cos_sim) 관계: # d^2 = 2(1 - cos_sim) # cos_sim = 1 - (d^2/2) cosine_similarity = 1 - (distance ** 2 / 2) # -1~1 범위의 코사인 유사도를 0~100 범위로 변환 similarity_score = ((cosine_similarity + 1) / 2) * 100 item_data = metadata.copy() item_data['similarity_score'] = similarity_score similar_items.append(item_data) similar_items.sort(key=lambda x: x['similarity_score'], reverse=True) return similar_items except Exception as e: logger.error(f"Multimodal search error: {e}") return [] def update_db_with_multimodal(): """DB를 멀티모달 방식으로 업데이트""" try: # 새 컬렉션 생성 collection = setup_multimodal_collection() # 기존 컬렉션에서 데이터 가져오기 client = chromadb.PersistentClient(path="./clothesDB_11GmarketMusinsa") old_collection = client.get_collection("clothes") old_data = old_collection.get(include=['metadatas']) total_items = len(old_data['metadatas']) progress_bar = st.progress(0) status_text = st.empty() batch_size = 100 successful_updates = 0 failed_updates = 0 for i in range(0, total_items, batch_size): batch = old_data['metadatas'][i:i + batch_size] images = [] valid_metadatas = [] valid_ids = [] for metadata in batch: try: if 'image_url' in metadata: response = requests.get(metadata['image_url']) img = Image.open(io.BytesIO(response.content)).convert('RGB') images.append(np.array(img)) valid_metadatas.append(metadata) valid_ids.append(metadata.get('id', str(hash(metadata['image_url'])))) successful_updates += 1 except Exception as e: logger.error(f"Error processing image: {e}") failed_updates += 1 continue if images: collection.add( ids=valid_ids, images=images, metadatas=valid_metadatas ) # Update progress progress = (i + len(batch)) / total_items progress_bar.progress(progress) status_text.text(f"Processing: {i + len(batch)}/{total_items} items. " f"Success: {successful_updates}, Failed: {failed_updates}") status_text.text(f"Update completed. Successfully processed: {successful_updates}, " f"Failed: {failed_updates}") return True except Exception as e: logger.error(f"Multimodal DB update error: {e}") return False def show_similar_items(similar_items): """Display similar items in a structured format with similarity scores""" if not similar_items: st.warning("No similar items found.") return st.subheader("Similar Items:") items_per_row = 2 for i in range(0, len(similar_items), items_per_row): cols = st.columns(items_per_row) for j, col in enumerate(cols): if i + j < len(similar_items): item = similar_items[i + j] with col: try: if 'image_url' in item: st.image(item['image_url'], use_column_width=True) st.markdown(f"**Similarity: {item['similarity_score']:.1f}%**") st.write(f"Brand: {item.get('brand', 'Unknown')}") name = item.get('name', 'Unknown') if len(name) > 50: name = name[:47] + "..." st.write(f"Name: {name}") price = item.get('price', 0) if isinstance(price, (int, float)): st.write(f"Price: {price:,}원") else: st.write(f"Price: {price}") if 'discount' in item and item['discount']: st.write(f"Discount: {item['discount']}%") if 'original_price' in item: st.write(f"Original: {item['original_price']:,}원") st.divider() except Exception as e: logger.error(f"Error displaying item: {e}") st.error("Error displaying this item") def process_search(image, mask, num_results): """유사 아이템 검색 처리""" try: with st.spinner("Finding similar items..."): similar_items = search_similar_items(image, mask, num_results) return similar_items except Exception as e: logger.error(f"Search processing error: {e}") return None def handle_file_upload(): if st.session_state.uploaded_file is not None: image = Image.open(st.session_state.uploaded_file).convert('RGB') st.session_state.image = image st.session_state.upload_state = 'image_uploaded' st.rerun() def handle_detection(): if st.session_state.image is not None: detected_items = process_segmentation(st.session_state.image) st.session_state.detected_items = detected_items st.session_state.upload_state = 'items_detected' st.rerun() def handle_search(): st.session_state.search_clicked = True def main(): st.title("Fashion Search App") # Admin controls in sidebar st.sidebar.title("Admin Controls") if st.sidebar.checkbox("Show Admin Interface"): if st.sidebar.button("Update Database (Multimodal)"): with st.spinner("Updating database with multimodal support..."): success = update_db_with_multimodal() if success: st.sidebar.success("Database updated successfully!") else: st.sidebar.error("Failed to update database") st.divider() # 파일 업로더 if st.session_state.upload_state == 'initial': uploaded_file = st.file_uploader("Upload an image", type=['png', 'jpg', 'jpeg'], key='uploaded_file', on_change=handle_file_upload) # 이미지가 업로드된 상태 if st.session_state.image is not None: st.image(st.session_state.image, caption="Uploaded Image", use_column_width=True) if st.session_state.detected_items is None: if st.button("Detect Items", key='detect_button', on_click=handle_detection): pass # 검출된 아이템 표시 및 검색 if st.session_state.detected_items is not None and len(st.session_state.detected_items) > 0: cols = st.columns(2) for idx, item in enumerate(st.session_state.detected_items): with cols[idx % 2]: try: if item.get('mask') is not None: masked_img = np.array(st.session_state.image) * np.expand_dims(item['mask'], axis=2) st.image(masked_img.astype(np.uint8), caption=f"Detected {item.get('label', 'Unknown')}") st.write(f"Item {idx + 1}: {item.get('label', 'Unknown')}") score = item.get('score') if score is not None and isinstance(score, (int, float)): st.write(f"Confidence: {score*100:.1f}%") else: st.write("Confidence: N/A") except Exception as e: logger.error(f"Error displaying item {idx}: {str(e)}") st.error(f"Error displaying item {idx}") valid_items = [i for i in range(len(st.session_state.detected_items)) if st.session_state.detected_items[i].get('mask') is not None] if not valid_items: st.warning("No valid items detected for search.") return selected_idx = st.selectbox( "Select item to search:", valid_items, format_func=lambda i: f"{st.session_state.detected_items[i].get('label', 'Unknown')}", key='item_selector' ) search_col1, search_col2 = st.columns([1, 2]) with search_col1: search_clicked = st.button("Search Similar Items", key='search_button', type="primary") with search_col2: num_results = st.slider("Number of results:", min_value=1, max_value=20, value=5, key='num_results') if search_clicked or st.session_state.get('search_clicked', False): st.session_state.search_clicked = True selected_item = st.session_state.detected_items[selected_idx] if selected_item.get('mask') is None: st.error("Selected item has no valid mask for search.") return if 'search_results' not in st.session_state: similar_items = process_search(st.session_state.image, selected_item['mask'], num_results) st.session_state.search_results = similar_items if st.session_state.search_results: show_similar_items(st.session_state.search_results) else: st.warning("No similar items found.") # 새 검색 버튼 if st.button("Start New Search", key='new_search'): for key in list(st.session_state.keys()): del st.session_state[key] st.rerun() if __name__ == "__main__": print('시작') main()