Update app.py

app.py

@@ -12,10 +12,6 @@ from inference_sdk import InferenceHTTPClient
 import matplotlib.pyplot as plt
 import base64
 
-# 오류 처리를 위한 예외 클래스 정의
-class APIError(Exception):
-    pass
-
 # Load model and tokenizer
 @st.cache_resource
 def load_model():
@@ -26,12 +22,62 @@ def load_model():
 
 model, preprocess_val, tokenizer, device = load_model()
 
-# Roboflow client setup function
-def setup_roboflow_client(api_key):
-    return InferenceHTTPClient(
-        api_url="https://outline.roboflow.com",
-        api_key=api_key
-    )
+# Load and process data
+@st.cache_data
+def load_data():
+    with open('musinsa-final.json', 'r', encoding='utf-8') as f:
+        return json.load(f)
+
+data = load_data()
+
+# Process database with segmentation
+@st.cache_data
+def process_database():
+    database_embeddings = []
+    database_info = []
+    for item in data:
+        image_url = item['이미지 링크'][0]
+        product_id = item.get('\ufeff상품 ID') or item.get('상품 ID')
+        
+        image = download_and_process_image(image_url)
+        if image is None:
+            continue
+        
+        # Save the image temporarily
+        temp_path = f"temp_{product_id}.jpg"
+        image.save(temp_path, 'JPEG')
+        
+        segmented_image = segment_image(temp_path)
+        embedding = get_image_embedding(segmented_image)
+        
+        database_embeddings.append(embedding)
+        database_info.append({
+            'id': product_id,
+            'category': item['카테고리'],
+            'brand': item['브랜드명'],
+            'name': item['제품명'],
+            'price': item['정가'],
+            'discount': item['할인율'],
+            'image_url': image_url
+        })
+    
+    return np.vstack(database_embeddings), database_info
+
+# Initialize database_embeddings and database_info
+database_embeddings, database_info = process_database()
+
+def find_similar_images(query_embedding, top_k=5):
+    similarities = np.dot(database_embeddings, query_embedding.T).squeeze()
+    top_indices = np.argsort(similarities)[::-1][:top_k]
+    
+    results = []
+    for idx in top_indices:
+        results.append({
+            'info': database_info[idx],
+            'similarity': similarities[idx]
+        })
+    
+    return results
 
 # Streamlit app
 st.title("Fashion Search App with Segmentation")
@@ -42,127 +88,6 @@ api_key = st.text_input("Enter your Roboflow API Key", type="password")
 if api_key:
     CLIENT = setup_roboflow_client(api_key)
 
-    def segment_image(image_path):
-        try:
-            # 이미지 파일 읽기
-            with open(image_path, "rb") as image_file:
-                image_data = image_file.read()
-            
-            # 이미지를 base64로 인코딩
-            encoded_image = base64.b64encode(image_data).decode('utf-8')
-            
-            # 원본 이미지 로드
-            image = cv2.imread(image_path)
-            image = cv2.resize(image, (800, 600))
-            mask = np.zeros(image.shape, dtype=np.uint8)
-            
-            try:
-                # Roboflow API 호출
-                results = CLIENT.infer(encoded_image, model_id="closet/1")
-            except Exception as api_error:
-                st.error(f"API Error: {str(api_error)}")
-                return Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
-            
-            if 'predictions' in results:
-                for prediction in results['predictions']:
-                    points = prediction['points']
-                    pts = np.array([[p['x'], p['y']] for p in points], np.int32)
-                    scale_x = image.shape[1] / results['image']['width']
-                    scale_y = image.shape[0] / results['image']['height']
-                    pts = pts * [scale_x, scale_y]
-                    pts = pts.astype(np.int32)
-                    pts = pts.reshape((-1, 1, 2))
-                    cv2.fillPoly(mask, [pts], color=(255, 255, 255))  # White mask
-                
-                segmented_image = cv2.bitwise_and(image, mask)
-            else:
-                st.warning("No predictions found in the image. Returning original image.")
-                segmented_image = image
-            
-            return Image.fromarray(cv2.cvtColor(segmented_image, cv2.COLOR_BGR2RGB))
-        except Exception as e:
-            st.error(f"Error in segmentation: {str(e)}")
-            # 원본 이미지를 다시 읽어 반환
-            return Image.open(image_path)
-    def get_image_embedding(image):
-        image_tensor = preprocess_val(image).unsqueeze(0).to(device)
-        with torch.no_grad():
-            image_features = model.encode_image(image_tensor)
-            image_features /= image_features.norm(dim=-1, keepdim=True)
-        return image_features.cpu().numpy()
-
-    # Load and process data
-    @st.cache_data
-    def load_data():
-        with open('musinsa-final.json', 'r', encoding='utf-8') as f:
-            return json.load(f)
-
-    data = load_data()
-
-    # Process database with segmentation
-    @st.cache_data
-    def download_and_process_image(image_url):
-        try:
-            response = requests.get(image_url)
-            response.raise_for_status()  # Raises an HTTPError for bad responses
-            image = Image.open(BytesIO(response.content))
-            
-            # Convert image to RGB mode if it's in RGBA mode
-            if image.mode == 'RGBA':
-                image = image.convert('RGB')
-            
-            return image
-        except requests.RequestException as e:
-            st.error(f"Error downloading image: {e}")
-            return None
-        except Exception as e:
-            st.error(f"Error processing image: {e}")
-            return None
-
-    def process_database():
-        database_embeddings = []
-        database_info = []
-        for item in data:
-            image_url = item['이미지 링크'][0]
-            product_id = item.get('\ufeff상품 ID') or item.get('상품 ID')
-            
-            image = download_and_process_image(image_url)
-            if image is None:
-                continue
-            
-            # Save the image temporarily
-            temp_path = f"temp_{product_id}.jpg"
-            image.save(temp_path, 'JPEG')
-            
-            segmented_image = segment_image(temp_path)
-            embedding = get_image_embedding(segmented_image)
-            
-            database_embeddings.append(embedding)
-            database_info.append({
-                'id': product_id,
-                'category': item['카테고리'],
-                'brand': item['브랜드명'],
-                'name': item['제품명'],
-                'price': item['정가'],
-                'discount': item['할인율'],
-                'image_url': image_url
-            })
-        
-        return np.vstack(database_embeddings), database_info
-
-    def find_similar_images(query_embedding, top_k=5):
-        similarities = np.dot(database_embeddings, query_embedding.T).squeeze()
-        top_indices = np.argsort(similarities)[::-1][:top_k]
-        
-        results = []
-        for idx in top_indices:
-            results.append({
-                'info': database_info[idx],
-                'similarity': similarities[idx]
-            })
-        
-        return results
-
     uploaded_file = st.file_uploader("Choose an image...", type="jpg")
     if uploaded_file is not None:
         image = Image.open(uploaded_file)
@@ -195,4 +120,101 @@ if api_key:
                         st.write(f"Discount: {img['info']['discount']}%")
                         st.write(f"Similarity: {img['similarity']:.2f}")
 else:
-    st.warning("Please enter your Roboflow API Key to use the app.")
+    st.warning("Please enter your Roboflow API Key to use the app.")
+
+
+# Process database with segmentation
+@st.cache_data
+def download_and_process_image(image_url):
+    try:
+        response = requests.get(image_url)
+        response.raise_for_status()  # Raises an HTTPError for bad responses
+        image = Image.open(BytesIO(response.content))
+        
+        # Convert image to RGB mode if it's in RGBA mode
+        if image.mode == 'RGBA':
+            image = image.convert('RGB')
+        
+        return image
+    except requests.RequestException as e:
+        st.error(f"Error downloading image: {e}")
+        return None
+    except Exception as e:
+        st.error(f"Error processing image: {e}")
+        return None
+
+def process_database():
+    database_embeddings = []
+    database_info = []
+    for item in data:
+        image_url = item['이미지 링크'][0]
+        product_id = item.get('\ufeff상품 ID') or item.get('상품 ID')
+        
+        image = download_and_process_image(image_url)
+        if image is None:
+            continue
+        
+        # Save the image temporarily
+        temp_path = f"temp_{product_id}.jpg"
+        image.save(temp_path, 'JPEG')
+        
+        segmented_image = segment_image(temp_path)
+        embedding = get_image_embedding(segmented_image)
+        
+        database_embeddings.append(embedding)
+        database_info.append({
+            'id': product_id,
+            'category': item['카테고리'],
+            'brand': item['브랜드명'],
+            'name': item['제품명'],
+            'price': item['정가'],
+            'discount': item['할인율'],
+            'image_url': image_url
+        })
+    
+    return np.vstack(database_embeddings), database_info
+
+def find_similar_images(query_embedding, top_k=5):
+    similarities = np.dot(database_embeddings, query_embedding.T).squeeze()
+    top_indices = np.argsort(similarities)[::-1][:top_k]
+    
+    results = []
+    for idx in top_indices:
+        results.append({
+            'info': database_info[idx],
+            'similarity': similarities[idx]
+        })
+    
+    return results
+
+uploaded_file = st.file_uploader("Choose an image...", type="jpg")
+if uploaded_file is not None:
+    image = Image.open(uploaded_file)
+    st.image(image, caption='Uploaded Image', use_column_width=True)
+    
+    if st.button('Find Similar Items'):
+        with st.spinner('Processing...'):
+            # Save uploaded image temporarily
+            temp_path = "temp_upload.jpg"
+            image.save(temp_path)
+            
+            # Segment the uploaded image
+            segmented_image = segment_image(temp_path)
+            st.image(segmented_image, caption='Segmented Image', use_column_width=True)
+            
+            # Get embedding for segmented image
+            query_embedding = get_image_embedding(segmented_image)
+            similar_images = find_similar_images(query_embedding)
+            
+            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']}")
+                    st.write(f"Category: {img['info']['category']}")
+                    st.write(f"Price: {img['info']['price']}")
+                    st.write(f"Discount: {img['info']['discount']}%")
+                    st.write(f"Similarity: {img['similarity']:.2f}")