Create app.py

app.py

@@ -0,0 +1,517 @@
+import streamlit as st
+import open_clip
+import torch
+from PIL import Image
+import numpy as np
+from transformers import pipeline
+import chromadb
+import logging
+import io
+import requests
+from concurrent.futures import ThreadPoolExecutor
+
+# 로깅 설정
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+# 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 models
+@st.cache_resource
+def load_models():
+    try:
+        # CLIP 모델
+        model, _, preprocess_val = open_clip.create_model_and_transforms('hf-hub:Marqo/marqo-fashionSigLIP')
+        
+        # 세그멘테이션 모델
+        segmenter = pipeline(model="mattmdjaga/segformer_b2_clothes")
+        
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        model.to(device)
+        
+        return model, preprocess_val, segmenter, device
+    except Exception as e:
+        logger.error(f"Error loading models: {e}")
+        raise
+
+# 모델 로드
+clip_model, preprocess_val, segmenter, device = load_models()
+
+# ChromaDB 설정
+client = chromadb.PersistentClient(path="./clothesDB_11GmarketMusinsa")
+collection = client.get_collection(name="clothes")
+
+def extract_color_histogram(image, mask=None):
+    """Extract color histogram from the image, considering the mask if provided"""
+    try:
+        img_array = np.array(image)
+        if mask is not None:
+            # Apply mask
+            mask = np.expand_dims(mask, axis=2)
+            img_array = img_array * mask
+            # Only consider pixels that are part of the clothing item
+            valid_pixels = img_array[mask[:,:,0] > 0]
+        else:
+            valid_pixels = img_array.reshape(-1, 3)
+            
+        # Convert to HSV color space for better color representation
+        if len(valid_pixels) > 0:
+            img_hsv = Image.fromarray(valid_pixels.reshape(1, -1, 3).astype(np.uint8)).convert('HSV')
+            hsv_pixels = np.array(img_hsv)
+            
+            # Calculate histogram for each HSV channel
+            h_hist = np.histogram(hsv_pixels[:,:,0], bins=10, range=(0, 256))[0]
+            s_hist = np.histogram(hsv_pixels[:,:,1], bins=10, range=(0, 256))[0]
+            v_hist = np.histogram(hsv_pixels[:,:,2], bins=10, range=(0, 256))[0]
+            
+            # Normalize histograms
+            h_hist = h_hist / h_hist.sum() if h_hist.sum() > 0 else h_hist
+            s_hist = s_hist / s_hist.sum() if s_hist.sum() > 0 else s_hist
+            v_hist = v_hist / v_hist.sum() if v_hist.sum() > 0 else v_hist
+            
+            return np.concatenate([h_hist, s_hist, v_hist])
+        return np.zeros(30)  # Return zero histogram if no valid pixels
+    except Exception as e:
+        logger.error(f"Color histogram extraction error: {e}")
+        return np.zeros(30)
+
+def process_segmentation(image):
+    """Segmentation processing"""
+    try:
+        # pipeline 출력 결과 직접 처리
+        output = segmenter(image)
+        
+        if not output or len(output) == 0:
+            logger.warning("No segments found in image")
+            return []
+            
+        processed_items = []
+        for segment in output:
+            # 기본값을 포함하여 딕셔너리 생성
+            processed_segment = {
+                'label': segment.get('label', 'Unknown'),
+                'score': segment.get('score', 1.0),  # score가 없으면 1.0을 기본값으로 사용
+                'mask': None
+            }
+            
+            mask = segment.get('mask')
+            if mask is not None:
+                # 마스크가 numpy array가 아닌 경우 변환
+                if not isinstance(mask, np.ndarray):
+                    mask = np.array(mask)
+                
+                # 마스크가 2D가 아닌 경우 첫 번째 채널 사용
+                if len(mask.shape) > 2:
+                    mask = mask[:, :, 0]
+                
+                # bool 마스크를 float로 변환
+                processed_segment['mask'] = mask.astype(float)
+            else:
+                logger.warning(f"No mask found for segment with label {processed_segment['label']}")
+                continue  # 마스크가 없는 세그먼트는 건너뜀
+            
+            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 extract_features(image, mask=None):
+    """Extract both CLIP features and color features with segmentation mask"""
+    try:
+        # Extract CLIP features
+        if mask is not None:
+            img_array = np.array(image)
+            mask = np.expand_dims(mask, axis=2)
+            masked_img = img_array * mask
+            masked_img[mask[:,:,0] == 0] = 255  # Set background to white
+            image = Image.fromarray(masked_img.astype(np.uint8))
+        
+        image_tensor = preprocess_val(image).unsqueeze(0).to(device)
+        with torch.no_grad():
+            clip_features = clip_model.encode_image(image_tensor)
+            clip_features /= clip_features.norm(dim=-1, keepdim=True)
+        clip_features = clip_features.cpu().numpy().flatten()
+        
+        # Extract color features
+        color_features = extract_color_histogram(image, mask)
+        
+        # Combine features
+        # Note: We normalize and weight the features to balance their influence
+        clip_features_normalized = clip_features / np.linalg.norm(clip_features)
+        color_features_normalized = color_features / np.linalg.norm(color_features)
+        
+        # Adjust these weights to control the influence of each feature type
+        clip_weight = 0.7  # CLIP features weight
+        color_weight = 0.3  # Color features weight
+        
+        combined_features = np.concatenate([
+            clip_features_normalized * clip_weight,
+            color_features_normalized * color_weight
+        ])
+        
+        return combined_features
+    except Exception as e:
+        logger.error(f"Feature extraction error: {e}")
+        raise
+
+def download_and_process_image(image_url, metadata_id):
+    """Download image from URL and apply segmentation"""
+    try:
+        response = requests.get(image_url, timeout=10)
+        if response.status_code != 200:
+            logger.error(f"Failed to download image {metadata_id}: HTTP {response.status_code}")
+            return None
+            
+        image = Image.open(io.BytesIO(response.content)).convert('RGB')
+        logger.info(f"Successfully downloaded image {metadata_id}")
+        
+        processed_items = process_segmentation(image)
+        if processed_items and len(processed_items) > 0:
+            # 가장 큰 세그먼트의 마스크 사용
+            largest_mask = max(processed_items, key=lambda x: np.sum(x['mask']))['mask']
+            features = extract_features(image, largest_mask)
+            logger.info(f"Successfully extracted features for image {metadata_id}")
+            return features
+            
+        logger.warning(f"No valid mask found for image {metadata_id}")
+        return None
+        
+    except Exception as e:
+        logger.error(f"Error processing image {metadata_id}: {str(e)}")
+        import traceback
+        logger.error(traceback.format_exc())
+        return None
+
+def update_db_with_segmentation():
+    """DB의 모든 이미지에 대해 segmentation을 적용하고 feature를 업데이트"""
+    try:
+        logger.info("Starting database update with segmentation and color features")
+        
+        # 새로운 collection 생성
+        try:
+            client.delete_collection("clothes_segmented")
+            logger.info("Deleted existing segmented collection")
+        except:
+            logger.info("No existing segmented collection to delete")
+            
+        new_collection = client.create_collection(
+            name="clothes_segmented",
+            metadata={"description": "Clothes collection with segmentation and color features"}
+        )
+        logger.info("Created new segmented collection")
+        
+        # 기존 collection에서 메타데이터만 가져오기
+        try:
+            all_items = collection.get(include=['metadatas'])
+            total_items = len(all_items['metadatas'])
+            logger.info(f"Found {total_items} items in database")
+        except Exception as e:
+            logger.error(f"Error getting items from collection: {str(e)}")
+            all_items = {'metadatas': []}
+            total_items = 0
+            
+        # 진행 상황 표시를 위한 progress bar
+        progress_bar = st.progress(0)
+        status_text = st.empty()
+        
+        successful_updates = 0
+        failed_updates = 0
+        
+        with ThreadPoolExecutor(max_workers=4) as executor:
+            futures = []
+            # 이미지 URL이 있는 항목만 처리
+            valid_items = [m for m in all_items['metadatas'] if 'image_url' in m]
+            
+            for metadata in valid_items:
+                future = executor.submit(
+                    download_and_process_image, 
+                    metadata['image_url'],
+                    metadata.get('id', 'unknown')
+                )
+                futures.append((metadata, future))
+            
+            # 결과 처리 및 새 DB에 저장
+            for idx, (metadata, future) in enumerate(futures):
+                try:
+                    new_features = future.result()
+                    if new_features is not None:
+                        item_id = metadata.get('id', str(hash(metadata['image_url'])))
+                        try:
+                            new_collection.add(
+                                embeddings=[new_features.tolist()],
+                                metadatas=[metadata],
+                                ids=[item_id]
+                            )
+                            successful_updates += 1
+                            logger.info(f"Successfully added item {item_id}")
+                        except Exception as e:
+                            logger.error(f"Error adding item to new collection: {str(e)}")
+                            failed_updates += 1
+                    else:
+                        failed_updates += 1
+                        
+                    # 진행 상황 업데이트
+                    progress = (idx + 1) / len(futures)
+                    progress_bar.progress(progress)
+                    status_text.text(f"Processing: {idx + 1}/{len(futures)} items. Success: {successful_updates}, Failed: {failed_updates}")
+                    
+                except Exception as e:
+                    logger.error(f"Error processing item: {str(e)}")
+                    failed_updates += 1
+                    continue
+        
+        # 최종 결과 표시
+        status_text.text(f"Update completed. Successfully processed: {successful_updates}, Failed: {failed_updates}")
+        logger.info(f"Database update completed. Successful: {successful_updates}, Failed: {failed_updates}")
+        
+        # 성공적으로 처리된 항목이 있는지 확인
+        if successful_updates > 0:
+            return True
+        else:
+            logger.error("No items were successfully processed")
+            return False
+        
+    except Exception as e:
+        logger.error(f"Database update error: {str(e)}")
+        import traceback
+        logger.error(traceback.format_exc())
+        return False
+
+def search_similar_items(features, top_k=10):
+    """Search similar items using combined features"""
+    try:
+        # 세그멘테이션이 적용된 collection이 있는지 확인
+        try:
+            search_collection = client.get_collection("clothes_segmented")
+            logger.info("Using segmented collection for search")
+        except:
+            # 없으면 기존 collection 사용
+            search_collection = collection
+            logger.info("Using original collection for search")
+        
+        results = search_collection.query(
+            query_embeddings=[features.tolist()],
+            n_results=top_k,
+            include=['metadatas', 'distances']
+        )
+        
+        if not results or not results['metadatas'] or not results['distances']:
+            logger.warning("No results returned from ChromaDB")
+            return []
+
+        similar_items = []
+        for metadata, distance in zip(results['metadatas'][0], results['distances'][0]):
+            try:
+                similarity_score = 1 / (1 + float(distance))
+                item_data = metadata.copy()
+                item_data['similarity_score'] = similarity_score
+                similar_items.append(item_data)
+            except Exception as e:
+                logger.error(f"Error processing search result: {str(e)}")
+                continue
+        
+        similar_items.sort(key=lambda x: x['similarity_score'], reverse=True)
+        return similar_items
+    except Exception as e:
+        logger.error(f"Search error: {str(e)}")
+        return []
+
+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:")
+    
+    # 결과를 2열로 표시
+    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)
+                        
+                        # 유사도 점수를 퍼센트로 표시
+                        similarity_percent = item['similarity_score'] * 100
+                        st.markdown(f"**Similarity: {similarity_percent:.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("Extracting features..."):
+            features = extract_features(image, mask)
+        
+        with st.spinner("Finding similar items..."):
+            similar_items = search_similar_items(features, top_k=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"):
+        # Admin interface 구현 (필요한 경우)
+        st.sidebar.warning("Admin interface is not implemented yet.")
+        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:
+            # 감지된 아이템들을 2열로 표시
+            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 값이 있고 숫자인 경우에만 표시
+                        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__":
+    main()