Upload 14 files

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+clothesDB_11GmarketMusinsa/chroma.sqlite3 filter=lfs diff=lfs merge=lfs -text

Dockerfile

@@ -0,0 +1,26 @@
+# Use Python 3.10
+FROM python:3.10-slim
+
+# Set the working directory in the container
+WORKDIR /app
+
+# Copy the current directory contents into the container at /app
+COPY . /app
+
+# Install system dependencies
+RUN apt-get update && apt-get install -y \
+    build-essential \
+    curl \
+    software-properties-common \
+    git \
+    && rm -rf /var/lib/apt/lists/*
+
+# Upgrade pip and install required python packages
+RUN pip install --no-cache-dir --upgrade pip
+RUN pip install --no-cache-dir -r requirements.txt
+
+# Make port 8501 available to the world outside this container
+EXPOSE 8501
+
+# Run app.py when the container launches
+CMD ["streamlit", "run", "app.py"]

README.md

@@ -1,13 +1,13 @@
----
-title: Itda Segment
-emoji: 🐨
-colorFrom: purple
-colorTo: blue
-sdk: streamlit
-sdk_version: 1.39.0
-app_file: app.py
-pinned: false
-license: apache-2.0
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+---
+title: Itda Nosegmentation
+emoji: 🐨
+colorFrom: gray
+colorTo: purple
+sdk: streamlit
+sdk_version: 1.38.0
+app_file: app.py
+pinned: false
+license: apache-2.0
+---
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,444 @@
+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 process_segmentation(image):
+    """Segmentation processing"""
+    try:
+        # pipeline 출력 결과 직접 처리
+        output = segmenter(image)
+        
+        if not output:
+            logger.warning("No segments found in image")
+            return None
+            
+        # 각 세그먼트의 마스크 크기 계산
+        segment_sizes = [np.sum(seg['mask']) for seg in output]
+        
+        if not segment_sizes:
+            return None
+            
+        # 가장 큰 세그먼트 선택
+        largest_idx = np.argmax(segment_sizes)
+        mask = output[largest_idx]['mask']
+        
+        # 마스크가 numpy array가 아닌 경우 변환
+        if not isinstance(mask, np.ndarray):
+            mask = np.array(mask)
+            
+        # 마스크가 2D가 아닌 경우 첫 번째 채널 사용
+        if len(mask.shape) > 2:
+            mask = mask[:, :, 0]
+            
+        # bool 마스크를 float로 변환
+        mask = mask.astype(float)
+        
+        logger.info(f"Successfully created mask with shape {mask.shape}")
+        return mask
+        
+    except Exception as e:
+        logger.error(f"Segmentation error: {str(e)}")
+        import traceback
+        logger.error(traceback.format_exc())
+        return None
+
+def download_and_process_image(image_url, metadata_id):
+    """Download image from URL and apply segmentation"""
+    try:
+        response = requests.get(image_url, timeout=10)  # timeout 추가
+        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}")
+        
+        mask = process_segmentation(image)
+        if mask is not None:
+            features = extract_features(image, 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")
+        
+        # 새로운 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-based 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 extract_features(image, mask=None):
+    """Extract CLIP features with segmentation mask"""
+    try:
+        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  # 배경을 흰색으로
+            image = Image.fromarray(masked_img.astype(np.uint8))
+        
+        image_tensor = preprocess_val(image).unsqueeze(0).to(device)
+        with torch.no_grad():
+            features = clip_model.encode_image(image_tensor)
+            features /= features.norm(dim=-1, keepdim=True)
+        return features.cpu().numpy().flatten()
+    except Exception as e:
+        logger.error(f"Feature extraction error: {e}")
+        raise
+
+def search_similar_items(features, top_k=10):
+    """Search similar items using segmentation-based 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
+
+# Callback functions
+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 admin_interface():
+    st.title("Admin Interface - DB Update")
+    if st.button("Update DB with Segmentation"):
+        with st.spinner("Updating database with segmentation... This may take a while..."):
+            success = update_db_with_segmentation()
+            if success:
+                st.success("Database successfully updated with segmentation-based features!")
+            else:
+                st.error("Failed to update database. Please check the logs.")
+
+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.divider()
+
+    # 파일 업로더 (upload_state가 initial일 때만 표시)
+    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:
+            # 감지된 아이템들을 2열로 표시
+            cols = st.columns(2)
+            for idx, item in enumerate(st.session_state.detected_items):
+                with cols[idx % 2]:
+                    mask = item['mask']
+                    masked_img = np.array(st.session_state.image) * np.expand_dims(mask, axis=2)
+                    st.image(masked_img.astype(np.uint8), caption=f"Detected {item['label']}")
+                    st.write(f"Item {idx + 1}: {item['label']}")
+                    st.write(f"Confidence: {item['score']*100:.1f}%")
+            
+            # 아이템 선택
+            selected_idx = st.selectbox(
+                "Select item to search:",
+                range(len(st.session_state.detected_items)),
+                format_func=lambda i: f"{st.session_state.detected_items[i]['label']}",
+                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_mask = st.session_state.detected_items[selected_idx]['mask']
+                
+                # 검색 결과를 세션 상태에 저장
+                if 'search_results' not in st.session_state:
+                    similar_items = process_search(st.session_state.image, selected_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()

app_10281200.py

@@ -0,0 +1,235 @@
+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
+
+# 로깅 설정
+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'
+
+# 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 process_segmentation(image):
+    """Segmentation processing 안녕하세요"""
+    try:
+        segments = segmenter(image)
+        valid_items = []
+        
+        for s in segments:
+            mask_array = np.array(s['mask'])
+            confidence = np.mean(mask_array)
+            
+            valid_items.append({
+                'score': confidence,
+                'label': s['label'],
+                'mask': mask_array
+            })
+        
+        return valid_items
+    except Exception as e:
+        logger.error(f"Segmentation error: {e}")
+        return []
+
+def extract_features(image, mask=None):
+    """Extract CLIP features"""
+    try:
+        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
+            image = Image.fromarray(masked_img.astype(np.uint8))
+        
+        image_tensor = preprocess_val(image).unsqueeze(0).to(device)
+        with torch.no_grad():
+            features = clip_model.encode_image(image_tensor)
+            features /= features.norm(dim=-1, keepdim=True)
+        return features.cpu().numpy().flatten()
+    except Exception as e:
+        logger.error(f"Feature extraction error: {e}")
+        raise
+
+def search_similar_items(features, top_k=10):
+    """Search similar items with distance scores"""
+    try:
+        results = collection.query(
+            query_embeddings=[features.tolist()],
+            n_results=top_k,
+            include=['metadatas', 'distances']  # distances 포함
+        )
+        
+        similar_items = []
+        for metadata, distance in zip(results['metadatas'][0], results['distances'][0]):
+            # 거리를 유사도 점수로 변환 (0~1 범위)
+            similarity_score = 1 / (1 + distance)
+            metadata['similarity_score'] = similarity_score  # 메타데이터에 점수 추가
+            similar_items.append(metadata)
+            
+        return similar_items
+    except Exception as e:
+        logger.error(f"Search error: {e}")
+        return []
+
+def show_similar_items(similar_items):
+    """Display similar items in a structured format with similarity scores"""
+    st.subheader("Similar Items:")
+    for item in similar_items:
+        col1, col2 = st.columns([1, 2])
+        with col1:
+            st.image(item['image_url'])
+        with col2:
+            # 유사도 점수를 퍼센트로 표시
+            similarity_percent = item['similarity_score'] * 100
+            st.write(f"Similarity: {similarity_percent:.1f}%")
+            st.write(f"Brand: {item.get('brand', 'Unknown')}")
+            st.write(f"Name: {item.get('name', 'Unknown')}")
+            st.write(f"Price: {item.get('price', 'Unknown'):,}원")
+            if 'discount' in item:
+                st.write(f"Discount: {item['discount']}%")
+                if 'original_price' in item:
+                    st.write(f"Original Price: {item['original_price']:,}원")
+
+# 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 
+
+def reset_state():
+    """Reset all session state variables"""
+    for key in list(st.session_state.keys()):
+        del st.session_state[key]
+
+# Callback functions
+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 demo!!!")
+
+    # 파일 업로더 (upload_state가 initial일 때만 표시)
+    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)
+
+    # 이미지가 업로드된 상태 df
+    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
+        
+        # 검출된 아이템 표시d
+        if st.session_state.detected_items:
+            # 감지된 아이템들d을 2열로 표시
+            cols = st.columns(2)
+            for idx, item in enumerate(st.session_state.detected_items):
+                with cols[idx % 2]:
+                    mask = item['mask']
+                    masked_img = np.array(st.session_state.image) * np.expand_dims(mask, axis=2)
+                    st.image(masked_img.astype(np.uint8), caption=f"Detected {item['label']}")
+                    st.write(f"Item {idx + 1}: {item['label']}")
+                    st.write(f"Confidence: {item['score']*100:.1f}%")
+            
+            # 아이템 선택
+            selected_idx = st.selectbox(
+                "Select item to search:",
+                range(len(st.session_state.detected_items)),
+                format_func=lambda i: f"{st.session_state.detected_items[i]['label']}",
+                key='item_selector'
+            )
+            st.session_state.selected_item_index = selected_idx
+            
+            # 유사 아이템 검색
+            col1, col2 = st.columns([1, 2])
+            with col1:
+                search_button = st.button("Search Similar Items", 
+                                        key='search_button', 
+                                        on_click=handle_search,
+                                        type="primary")  # 강조된 버튼
+            with col2:
+                num_results = st.slider("Number of results:", 
+                                      min_value=1, 
+                                      max_value=20, 
+                                      value=5,
+                                      key='num_results')
+            
+            if st.session_state.search_clicked:
+                with st.spinner("Searching similar items..."):
+                    try:
+                        selected_mask = st.session_state.detected_items[selected_idx]['mask']
+                        features = extract_features(st.session_state.image, selected_mask)
+                        similar_items = search_similar_items(features, top_k=num_results)
+                        
+                        if similar_items:
+                            show_similar_items(similar_items)
+                        else:
+                            st.warning("No similar items found.")
+                    except Exception as e:
+                        st.error(f"Error during search: {str(e)}")
+
+    # 새 검색 버튼
+    if st.button("Start New Search  ", key='new_search'):
+        reset_state()
+        st.rerun()
+
+if __name__ == "__main__":
+    main()

app_accessary.py

@@ -0,0 +1,216 @@
+import streamlit as st
+import open_clip
+import torch
+import requests
+from PIL import Image
+from io import BytesIO
+import time
+import json
+import numpy as np
+import onnxruntime as ort
+from ultralytics import YOLO
+import cv2
+import chromadb
+
+@st.cache_resource
+def load_clip_model():
+    model, preprocess_train, preprocess_val = open_clip.create_model_and_transforms('hf-hub:Marqo/marqo-fashionSigLIP')
+    tokenizer = open_clip.get_tokenizer('hf-hub:Marqo/marqo-fashionSigLIP')
+    
+    # 파인튜닝한 모델의 state_dict 불러오기
+    #state_dict = torch.load('./accessory_clip.pt', map_location=torch.device('cpu'))
+    #model.load_state_dict(state_dict)  # 모델에 state_dict 적용
+    
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model.to(device)
+    
+    return model, preprocess_val, tokenizer, device
+
+clip_model, preprocess_val, tokenizer, device = load_clip_model()
+
+
+@st.cache_resource
+def load_yolo_model():
+    return YOLO("./accessaries.pt")
+
+yolo_model = load_yolo_model()
+
+# URL에서 이미지 로드
+def load_image_from_url(url, max_retries=3):
+    for attempt in range(max_retries):
+        try:
+            response = requests.get(url, timeout=10)
+            response.raise_for_status()
+            img = Image.open(BytesIO(response.content)).convert('RGB')
+            return img
+        except (requests.RequestException, Image.UnidentifiedImageError) as e:
+            if attempt < max_retries - 1:
+                time.sleep(1)
+            else:
+                return None
+
+# ChromaDB 클라이언트 설정
+client = chromadb.PersistentClient(path="./accessaryDB")
+collection = client.get_collection(name="accessary_items_ver2")
+
+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()
+
+def get_text_embedding(text):
+    text_tokens = tokenizer([text]).to(device)
+    with torch.no_grad():
+        text_features = clip_model.encode_text(text_tokens)
+        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)
+
+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()
+    top_indices = np.argsort(similarities)[::-1][:top_k]
+    
+    all_data = collection.get(include=['metadatas'])['metadatas']
+    
+    top_metadatas = [all_data[idx] for idx in top_indices]
+    
+    results = []
+    for idx, metadata in enumerate(top_metadatas):
+        results.append({
+            'info': metadata,
+            'similarity': similarities[top_indices[idx]]
+        })
+    return results
+
+def detect_clothing(image):
+    results = yolo_model(image)
+    detections = results[0].boxes.data.cpu().numpy()
+    categories = []
+    for detection in detections:
+        x1, y1, x2, y2, conf, cls = detection
+        category = yolo_model.names[int(cls)]
+        if category in ['Bracelets', 'Broches', 'bag', 'belt', 'earring', 'maangtika', 'necklace', 'nose ring', 'ring', 'tiara']:
+            categories.append({
+                'category': category,
+                'bbox': [int(x1), int(y1), int(x2), int(y2)],
+                'confidence': conf
+            })
+    return categories
+
+# 이미지 자르기
+def crop_image(image, bbox):
+    return image.crop((bbox[0], bbox[1], bbox[2], bbox[3]))
+
+# 세션 상태 초기화
+if 'step' not in st.session_state:
+    st.session_state.step = 'input'
+if 'query_image_url' not in st.session_state:
+    st.session_state.query_image_url = ''
+if 'detections' not in st.session_state:
+    st.session_state.detections = []
+if 'selected_category' not in st.session_state:
+    st.session_state.selected_category = None
+
+# Streamlit app
+st.title("Accessary Search App")
+
+# 단계별 처리
+if st.session_state.step == 'input':
+    st.session_state.query_image_url = st.text_input("Enter image URL:", st.session_state.query_image_url)
+    if st.button("Detect accesseary"):
+        if st.session_state.query_image_url:
+            query_image = load_image_from_url(st.session_state.query_image_url)
+            if query_image is not None:
+                st.session_state.query_image = query_image
+                st.session_state.detections = detect_clothing(query_image)
+                if st.session_state.detections:
+                    st.session_state.step = 'select_category'
+                else:
+                    st.warning("No items detected in the image.")
+            else:
+                st.error("Failed to load the image. Please try another URL.")
+        else:
+            st.warning("Please enter an image URL.")
+
+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:")
+    
+    for detection in st.session_state.detections:
+        col1, col2 = st.columns([1, 3])
+        with col1:
+            st.write(f"{detection['category']} (Confidence: {detection['confidence']:.2f})")
+        with col2:
+            cropped_image = crop_image(st.session_state.query_image, detection['bbox'])
+            st.image(cropped_image, caption=detection['category'], use_column_width=True)
+    
+    options = [f"{d['category']} (Confidence: {d['confidence']:.2f})" for d in st.session_state.detections]
+    selected_option = st.selectbox("Select a category to search:", options)
+    
+    if st.button("Search Similar Items"):
+        st.session_state.selected_category = selected_option
+        st.session_state.step = 'show_results'
+
+elif st.session_state.step == 'show_results':
+    st.image(st.session_state.query_image, caption="Query Image", use_column_width=True)
+    selected_detection = next(d for d in st.session_state.detections 
+                              if f"{d['category']} (Confidence: {d['confidence']:.2f})" == st.session_state.selected_category)
+    cropped_image = crop_image(st.session_state.query_image, selected_detection['bbox'])
+    st.image(cropped_image, caption="Cropped Image", use_column_width=True)
+    query_embedding = get_image_embedding(cropped_image)
+    similar_images = find_similar_images(query_embedding, collection)
+    
+    st.subheader("Similar Items:")
+    for img in similar_images:
+        col1, col2 = st.columns(2)
+        with col1:
+            st.image(img['info']['image_url'], use_column_width=True)
+        with col2:
+            st.write(f"Name: {img['info']['name']}")
+            st.write(f"Brand: {img['info']['brand']}")
+            category = img['info'].get('category')
+            if category:  
+                st.write(f"Category: {category}")
+            st.write(f"Price: {img['info']['price']}")
+            st.write(f"Discount: {img['info']['discount']}%")
+            st.write(f"Similarity: {img['similarity']:.2f}")
+    
+    if st.button("Start New Search"):
+        st.session_state.step = 'input'
+        st.session_state.query_image_url = ''
+        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"):
+        if query_text:
+            text_embedding = get_text_embedding(query_text)
+            similar_images = find_similar_images(text_embedding, collection)
+            st.subheader("Similar Items:")
+            for img in similar_images:
+                col1, col2 = st.columns(2)
+                with col1:
+                    st.image(img['info']['image_url'], use_column_width=True)
+                with col2:
+                    st.write(f"Name: {img['info']['name']}")
+                    st.write(f"Brand: {img['info']['brand']}")
+                    category = img['info'].get('category')
+                    if category:  
+                        st.write(f"Category: {category}")
+                    st.write(f"Price: {img['info']['price']}")
+                    st.write(f"Discount: {img['info']['discount']}%")
+                    st.write(f"Similarity: {img['similarity']:.2f}")
+        else:
+            st.warning("Please enter a search text.")

clothesDB_11GmarketMusinsa/b4c365a0-3372-4031-9ff8-d457b4123d0a/data_level0.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:71986228e55386105f7be1ee4986b077f01efe1d3dcfc901cc1e4a61138af1dc
+size 3212000

clothesDB_11GmarketMusinsa/b4c365a0-3372-4031-9ff8-d457b4123d0a/header.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2cbe47b73ab24fff8fda6ff745cf0c3153d60098f2b086c4e2f11f4dd46c39a9
+size 100

clothesDB_11GmarketMusinsa/b4c365a0-3372-4031-9ff8-d457b4123d0a/index_metadata.pickle

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ab67b060118c79acfd5c73f1014457c1edf4745681f9bed95f08a422bd820656
+size 346027

clothesDB_11GmarketMusinsa/b4c365a0-3372-4031-9ff8-d457b4123d0a/length.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4333c839a3d07797dedf8ee38d3a56f0ad5e42e45babb1bf94d830617d043fcb
+size 4000

clothesDB_11GmarketMusinsa/b4c365a0-3372-4031-9ff8-d457b4123d0a/link_lists.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0ba1d7a2bccf360384cbf836c3cd29510268a93a83e100577fca177d77433f64
+size 8420

clothesDB_11GmarketMusinsa/chroma.sqlite3

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:509eb23fbf078e21b37f013125481a4bc07f9a34a90a69a2930bb9c0784aabba
+size 17100800

db_creation.log

@@ -0,0 +1,2 @@
+2024-10-28 13:02:14,164 - INFO - Loading models...
+2024-10-28 13:02:14,166 - INFO - Loading models...

db_segmentation.py

@@ -0,0 +1,221 @@
+import chromadb
+import logging
+import open_clip
+import torch
+from PIL import Image
+import numpy as np
+from transformers import pipeline
+import requests
+import io
+from concurrent.futures import ThreadPoolExecutor
+from tqdm import tqdm
+import os
+
+# 로깅 설정
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(levelname)s - %(message)s',
+    handlers=[
+        logging.FileHandler('db_creation.log'),
+        logging.StreamHandler()
+    ]
+)
+logger = logging.getLogger(__name__)
+
+def load_models():
+    """Load CLIP and segmentation models"""
+    try:
+        logger.info("Loading models...")
+        # 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")
+        logger.info(f"Using device: {device}")
+        model.to(device)
+        
+        return model, preprocess_val, segmenter, device
+    except Exception as e:
+        logger.error(f"Error loading models: {e}")
+        raise
+
+def process_segmentation(image, segmenter):
+    """Apply segmentation to image"""
+    try:
+        segments = segmenter(image)
+        if not segments:
+            return None
+            
+        # 가장 큰 세그먼트 선택
+        largest_segment = max(segments, key=lambda s: np.sum(s['mask']))
+        mask = np.array(largest_segment['mask'])
+        
+        return mask
+        
+    except Exception as e:
+        logger.error(f"Segmentation error: {e}")
+        return None
+
+def extract_features(image, mask, model, preprocess_val, device):
+    """Extract CLIP features with segmentation mask"""
+    try:
+        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  # 배경을 흰색으로
+            image = Image.fromarray(masked_img.astype(np.uint8))
+        
+        image_tensor = preprocess_val(image).unsqueeze(0).to(device)
+        with torch.no_grad():
+            features = model.encode_image(image_tensor)
+            features /= features.norm(dim=-1, keepdim=True)
+        return features.cpu().numpy().flatten()
+    except Exception as e:
+        logger.error(f"Feature extraction error: {e}")
+        return None
+
+def download_and_process_image(url, metadata_id, model, preprocess_val, segmenter, device):
+    """Download and process single image"""
+    try:
+        response = requests.get(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')
+        
+        # Apply segmentation
+        mask = process_segmentation(image, segmenter)
+        if mask is None:
+            logger.warning(f"No valid mask found for image {metadata_id}")
+            return None
+            
+        # Extract features
+        features = extract_features(image, mask, model, preprocess_val, device)
+        if features is None:
+            logger.warning(f"Failed to extract features for image {metadata_id}")
+            return None
+            
+        return features
+        
+    except Exception as e:
+        logger.error(f"Error processing image {metadata_id}: {e}")
+        return None
+
+def create_segmented_db(source_path, target_path, batch_size=100):
+    """Create new segmented database from existing one"""
+    try:
+        logger.info("Loading models...")
+        model, preprocess_val, segmenter, device = load_models()
+        
+        # Source DB 연결
+        source_client = chromadb.PersistentClient(path=source_path)
+        source_collection = source_client.get_collection(name="clothes")
+        
+        # Target DB 생성
+        os.makedirs(target_path, exist_ok=True)
+        target_client = chromadb.PersistentClient(path=target_path)
+        
+        try:
+            target_client.delete_collection("clothes_segmented")
+        except:
+            pass
+            
+        target_collection = target_client.create_collection(
+            name="clothes_segmented",
+            metadata={"description": "Clothes collection with segmentation-based features"}
+        )
+        
+        # 전체 아이템 수 확인
+        all_items = source_collection.get(include=['metadatas'])
+        total_items = len(all_items['metadatas'])
+        logger.info(f"Found {total_items} items in source database")
+        
+        # 배치 처리를 위한 준비
+        successful_updates = 0
+        failed_updates = 0
+        
+        # ThreadPoolExecutor ���정
+        max_workers = min(10, os.cpu_count() or 4)
+        
+        with ThreadPoolExecutor(max_workers=max_workers) as executor:
+            # 전체 데이터를 배치로 나누어 처리
+            for batch_start in tqdm(range(0, total_items, batch_size), desc="Processing batches"):
+                batch_end = min(batch_start + batch_size, total_items)
+                batch_items = all_items['metadatas'][batch_start:batch_end]
+                
+                # 배치 내의 모든 이미지에 대한 future 생성
+                futures = []
+                for metadata in batch_items:
+                    if 'image_url' in metadata:
+                        future = executor.submit(
+                            download_and_process_image,
+                            metadata['image_url'],
+                            metadata.get('id', 'unknown'),
+                            model, preprocess_val, segmenter, device
+                        )
+                        futures.append((metadata, future))
+                
+                # 배치 결과 처리
+                batch_embeddings = []
+                batch_metadatas = []
+                batch_ids = []
+                
+                for metadata, future in futures:
+                    try:
+                        features = future.result()
+                        if features is not None:
+                            batch_embeddings.append(features.tolist())
+                            batch_metadatas.append(metadata)
+                            batch_ids.append(metadata.get('id', str(hash(metadata['image_url']))))
+                            successful_updates += 1
+                        else:
+                            failed_updates += 1
+                            
+                    except Exception as e:
+                        logger.error(f"Error processing batch item: {e}")
+                        failed_updates += 1
+                        continue
+                
+                # 배치 데이터 저장
+                if batch_embeddings:
+                    try:
+                        target_collection.add(
+                            embeddings=batch_embeddings,
+                            metadatas=batch_metadatas,
+                            ids=batch_ids
+                        )
+                        logger.info(f"Added batch of {len(batch_embeddings)} items")
+                    except Exception as e:
+                        logger.error(f"Error adding batch to collection: {e}")
+                        failed_updates += len(batch_embeddings)
+                        successful_updates -= len(batch_embeddings)
+        
+        # 최종 결과 출력
+        logger.info(f"Database creation completed.")
+        logger.info(f"Successfully processed: {successful_updates}")
+        logger.info(f"Failed: {failed_updates}")
+        logger.info(f"Total completion rate: {(successful_updates/total_items)*100:.2f}%")
+        
+        return True
+        
+    except Exception as e:
+        logger.error(f"Database creation error: {e}")
+        return False
+
+if __name__ == "__main__":
+    # 설정값
+    SOURCE_DB_PATH = "./clothesDB_11GmarketMusinsa"  # 원본 DB 경로
+    TARGET_DB_PATH = "./clothesDB_11GmarketMusinsa_segmented"  # 새로운 DB 경로
+    BATCH_SIZE = 50  # 한 번에 처리할 아이템 수
+    
+    # DB 생성 실행
+    success = create_segmented_db(SOURCE_DB_PATH, TARGET_DB_PATH, BATCH_SIZE)
+    
+    if success:
+        logger.info("Successfully created segmented database!")
+    else:
+        logger.error("Failed to create segmented database.")