Update app.py

app.py

@@ -63,23 +63,30 @@ def process_segmentation(image):
             
         processed_items = []
         for segment in output:
-            mask = segment['mask']
-            # 마스크가 numpy array가 아닌 경우 변환
-            if not isinstance(mask, np.ndarray):
-                mask = np.array(mask)
+            # 기본값을 포함하여 딕셔너리 생성
+            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]
+                # 마스크가 2D가 아닌 경우 첫 번째 채널 사용
+                if len(mask.shape) > 2:
+                    mask = mask[:, :, 0]
                 
-            # bool 마스크를 float로 변환
-            mask = mask.astype(float)
+                # 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({
-                'mask': mask,
-                'label': segment.get('label', 'Unknown'),
-                'score': segment.get('score', 0.0)
-            })
+            processed_items.append(processed_segment)
             
         logger.info(f"Successfully processed {len(processed_items)} segments")
         return processed_items
@@ -380,17 +387,35 @@ def main():
             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}%")
+                    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:",
-                range(len(st.session_state.detected_items)),
-                format_func=lambda i: f"{st.session_state.detected_items[i]['label']}",
+                valid_items,
+                format_func=lambda i: f"{st.session_state.detected_items[i].get('label', 'Unknown')}",
                 key='item_selector'
             )
             
@@ -410,11 +435,15 @@ def main():
             # 검색 결과 처리
             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']
+                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_mask, num_results)
+                    similar_items = process_search(st.session_state.image, selected_item['mask'], num_results)
                     st.session_state.search_results = similar_items
                 
                 # 저장된 검색 결과 표시