Gradio UI에 벡터 DB 저장 및 검색 기능 통합

app.py

@@ -69,7 +69,7 @@ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 print(f"Using device: {device}")
 
 # 벡터 DB에 객체 저장 함수
-def save_objects_to_vector_db(image, detection_results):
+def save_objects_to_vector_db(image, detection_results, model_type='yolo'):
     if image is None or detection_results is None:
         return "이미지나 객체 인식 결과가 없습니다."
     
@@ -79,36 +79,126 @@ def save_objects_to_vector_db(image, detection_results):
         image.save(buffered, format="JPEG")
         img_str = base64.b64encode(buffered.getvalue()).decode('utf-8')
         
-        # 객체 정보 추출
-        objects = []
-        for obj in detection_results['objects']:
-            objects.append({
-                "class": obj['class'],
-                "confidence": obj['confidence'],
-                "bbox": obj['bbox']
-            })
+        # 모델 타입에 따라 다른 API 엔드포인트 호출
+        if model_type in ['yolo', 'detr']:
+            # 객체 정보 추출
+            objects = []
+            for obj in detection_results['objects']:
+                objects.append({
+                    "class": obj['class'],
+                    "confidence": obj['confidence'],
+                    "bbox": obj['bbox']
+                })
+            
+            # API 요청 데이터 구성
+            data = {
+                "image": img_str,
+                "objects": objects,
+                "image_id": str(uuid.uuid4())
+            }
+            
+            # API 호출
+            response = requests.post(
+                "http://localhost:7860/api/add-detected-objects", 
+                json=data
+            )
+            
+            if response.status_code == 200:
+                result = response.json()
+                if 'error' in result:
+                    return f"오류 발생: {result['error']}"
+                return f"벡터 DB에 {len(objects)}개 객체 저장 완료! ID: {result.get('ids', '알 수 없음')}"
         
-        # API 요청 데이터 구성
-        data = {
-            "image": img_str,
-            "objects": objects,
-            "image_id": str(uuid.uuid4())
-        }
+        elif model_type == 'vit':
+            # ViT 분류 결과 저장
+            data = {
+                "image": img_str,
+                "metadata": {
+                    "model": "vit",
+                    "classifications": detection_results.get('classifications', [])
+                }
+            }
+            
+            # API 호출
+            response = requests.post(
+                "http://localhost:7860/api/add-image", 
+                json=data
+            )
+            
+            if response.status_code == 200:
+                result = response.json()
+                if 'error' in result:
+                    return f"오류 발생: {result['error']}"
+                return f"벡터 DB에 이미지 및 분류 결과 저장 완료! ID: {result.get('id', '알 수 없음')}"
+        
+        else:
+            return "지원하지 않는 모델 타입입니다."
+            
+        if response.status_code != 200:
+            return f"API 오류: {response.status_code}"
+    except Exception as e:
+        return f"오류 발생: {str(e)}"
+
+# 벡터 DB에서 유사 객체 검색 함수
+def search_similar_objects(image=None, class_name=None):
+    try:
+        data = {}
+        
+        if image is not None:
+            # 이미지를 base64로 인코딩
+            buffered = BytesIO()
+            image.save(buffered, format="JPEG")
+            img_str = base64.b64encode(buffered.getvalue()).decode('utf-8')
+            data["image"] = img_str
+            data["n_results"] = 5
+        elif class_name is not None and class_name.strip():
+            data["class_name"] = class_name.strip()
+            data["n_results"] = 5
+        else:
+            return "이미지나 클래스 이름 중 하나는 제공해야 합니다.", []
         
         # API 호출
         response = requests.post(
-            "http://localhost:7860/api/add-detected-objects", 
+            "http://localhost:7860/api/search-similar-objects", 
             json=data
         )
         
         if response.status_code == 200:
-            result = response.json()
-            return f"벡터 DB에 {len(objects)}개 객체 저장 성공! 저장된 객체 ID: {', '.join(result.get('object_ids', [])[:3])}..."
+            results = response.json()
+            if isinstance(results, dict) and 'error' in results:
+                return f"오류 발생: {results['error']}", []
+                
+            # 결과 포맷팅
+            formatted_results = []
+            for i, result in enumerate(results):
+                similarity = (1 - result.get('distance', 0)) * 100
+                img_data = result.get('image', '')
+                
+                # 이미지 데이터를 PIL 이미지로 변환
+                if img_data:
+                    try:
+                        img_bytes = base64.b64decode(img_data)
+                        img = Image.open(BytesIO(img_bytes))
+                    except Exception:
+                        img = None
+                else:
+                    img = None
+                
+                # 메타데이터 추출
+                metadata = result.get('metadata', {})
+                class_name = metadata.get('class', 'N/A')
+                confidence = metadata.get('confidence', 0) * 100 if metadata.get('confidence') else 'N/A'
+                
+                formatted_results.append({
+                    'image': img,
+                    'info': f"결과 #{i+1} | 유사도: {similarity:.2f}% | 클래스: {class_name} | 신뢰도: {confidence if isinstance(confidence, str) else f'{confidence:.2f}%'} | ID: {result.get('id', 'N/A')}"
+                })
+                
+            return f"{len(formatted_results)}개의 유사 객체를 찾았습니다.", formatted_results
         else:
-            return f"저장 실패: {response.text}"
-    
+            return f"API 오류: {response.status_code}", []
     except Exception as e:
-        return f"오류 발생: {str(e)}"
+        return f"오류 발생: {str(e)}", []
 
 # Define model inference functions
 def process_yolo(image):
@@ -334,36 +424,133 @@ with gr.Blocks(title="Object Detection Demo") as demo:
     # 벡터 DB 저장 버튼 및 결과 표시
     with gr.Row():
         with gr.Column():
-            save_to_db_button = gr.Button("YOLOv8 인식 결과를 벡터 DB에 저장", variant="primary")
+            gr.Markdown("### 벡터 DB 저장")
+            save_yolo_button = gr.Button("YOLOv8 인식 결과 저장", variant="primary")
+            save_detr_button = gr.Button("DETR 인식 결과 저장", variant="primary")
+            save_vit_button = gr.Button("ViT 분류 결과 저장", variant="primary")
             save_result = gr.Textbox(label="벡터 DB 저장 결과")
+        
+        with gr.Column():
+            gr.Markdown("### 벡터 DB 검색")
+            search_class = gr.Textbox(label="클래스 이름으로 검색")
+            search_button = gr.Button("검색", variant="secondary")
+            search_result_text = gr.Textbox(label="검색 결과 정보")
+            search_result_gallery = gr.Gallery(label="검색 결과", columns=5, height=300)
     
     # 객체 인식 결과 저장용 상태 변수
-    detection_state = gr.State(None)
+    yolo_detection_state = gr.State(None)
+    detr_detection_state = gr.State(None)
+    vit_classification_state = gr.State(None)
     
     # Set up event handlers
-    yolo_result = yolo_button.click(
+    yolo_button.click(
         fn=process_yolo,
         inputs=input_image,
-        outputs=[yolo_output, yolo_text, detection_state]
+        outputs=[yolo_output, yolo_text, yolo_detection_state]
     )
     
-    # 벡터 DB 저장 버튼 이벤트 핸들러
-    save_to_db_button.click(
-        fn=save_objects_to_vector_db,
-        inputs=[input_image, detection_state],
-        outputs=save_result
-    )
+    # DETR 결과 처리 함수 수정 - 상태 저장 추가
+    def process_detr_with_state(image):
+        result_image, result_text = process_detr(image)
+        
+        # 객체 인식 결과 추출
+        detection_results = {"objects": []}
+        
+        # 결과 텍스트에서 객체 정보 추출
+        lines = result_text.split('\n')
+        for line in lines:
+            if ': ' in line and ' at ' in line:
+                try:
+                    class_conf, location = line.split(' at ')
+                    class_name, confidence = class_conf.split(': ')
+                    confidence = float(confidence)
+                    
+                    # 바운딩 박스 정보 추출
+                    bbox_str = location.strip('[]').split(', ')
+                    bbox = [int(coord) for coord in bbox_str]
+                    
+                    detection_results["objects"].append({
+                        "class": class_name,
+                        "confidence": confidence,
+                        "bbox": bbox
+                    })
+                except Exception:
+                    pass
+        
+        return result_image, result_text, detection_results
+    
+    # ViT 결과 처리 함수 수정 - 상태 저장 추가
+    def process_vit_with_state(image):
+        result_text = process_vit(image)
+        
+        # 분류 결과 추출
+        classifications = []
+        
+        # 결과 텍스트에서 분류 정보 추출
+        lines = result_text.split('\n')
+        for line in lines:
+            if '. ' in line and ': ' in line:
+                try:
+                    rank_class, confidence = line.split(': ')
+                    _, class_name = rank_class.split('. ')
+                    confidence = float(confidence)
+                    
+                    classifications.append({
+                        "class": class_name,
+                        "confidence": confidence
+                    })
+                except Exception:
+                    pass
+        
+        return result_text, {"classifications": classifications}
     
     detr_button.click(
-        fn=process_detr,
+        fn=process_detr_with_state,
         inputs=input_image,
-        outputs=[detr_output, detr_text]
+        outputs=[detr_output, detr_text, detr_detection_state]
     )
     
     vit_button.click(
-        fn=process_vit,
+        fn=process_vit_with_state,
         inputs=input_image,
-        outputs=vit_text
+        outputs=[vit_text, vit_classification_state]
+    )
+    
+    # 벡터 DB 저장 버튼 이벤트 핸들러
+    save_yolo_button.click(
+        fn=lambda img, det: save_objects_to_vector_db(img, det, 'yolo'),
+        inputs=[input_image, yolo_detection_state],
+        outputs=save_result
+    )
+    
+    save_detr_button.click(
+        fn=lambda img, det: save_objects_to_vector_db(img, det, 'detr'),
+        inputs=[input_image, detr_detection_state],
+        outputs=save_result
+    )
+    
+    save_vit_button.click(
+        fn=lambda img, det: save_objects_to_vector_db(img, det, 'vit'),
+        inputs=[input_image, vit_classification_state],
+        outputs=save_result
+    )
+    
+    # 검색 버튼 이벤트 핸들러
+    def format_search_results(result_text, results):
+        images = []
+        captions = []
+        
+        for result in results:
+            if result.get('image'):
+                images.append(result['image'])
+                captions.append(result['info'])
+        
+        return result_text, [(img, cap) for img, cap in zip(images, captions)]
+    
+    search_button.click(
+        fn=lambda class_name: search_similar_objects(class_name=class_name),
+        inputs=search_class,
+        outputs=[search_result_text, search_result_gallery]
     )