Add vector DB save feature to Gradio UI for object detection results

app.py

@@ -3,6 +3,11 @@ import torch
 from PIL import Image
 import numpy as np
 import os
+import requests
+import json
+import base64
+from io import BytesIO
+import uuid
 
 # Model initialization
 print("Loading models... This may take a moment.")
@@ -63,10 +68,52 @@ import torch
 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):
+    if image is None or detection_results is None:
+        return "이미지나 객체 인식 결과가 없습니다."
+    
+    try:
+        # 이미지를 base64로 인코딩
+        buffered = BytesIO()
+        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 요청 데이터 구성
+        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()
+            return f"벡터 DB에 {len(objects)}개 객체 저장 성공! 저장된 객체 ID: {', '.join(result.get('object_ids', [])[:3])}..."
+        else:
+            return f"저장 실패: {response.text}"
+    
+    except Exception as e:
+        return f"오류 발생: {str(e)}"
+
 # Define model inference functions
 def process_yolo(image):
     if yolo_model is None:
-        return None, "YOLOv8 model not loaded"
+        return None, "YOLOv8 model not loaded", None
     
     # Measure inference time
     import time
@@ -91,13 +138,23 @@ def process_yolo(image):
     
     # Format detection results
     detections = []
+    detection_objects = {'objects': []}
+    
     for box in boxes:
         class_id = int(box.cls[0].item())
         class_name = class_names[class_id]
         confidence = round(box.conf[0].item(), 2)
         bbox = box.xyxy[0].tolist()
         bbox = [round(x) for x in bbox]
+        
         detections.append("{}: {} at {}".format(class_name, confidence, bbox))
+        
+        # 벡터 DB 저장용 객체 정보 추가
+        detection_objects['objects'].append({
+            'class': class_name,
+            'confidence': confidence,
+            'bbox': bbox
+        })
     
     # Calculate inference time
     inference_time = time.time() - start_time
@@ -108,6 +165,8 @@ def process_yolo(image):
     detection_text = "\n".join(detections) if detections else "No objects detected"
     detection_text += performance_info
     
+    return result_image, detection_text, detection_objects
+    
     return result_image, detection_text
 
 def process_detr(image):
@@ -272,11 +331,27 @@ with gr.Blocks(title="Object Detection Demo") as demo:
         with gr.Column():
             vit_text = gr.Textbox(label="ViT Classification Results")
     
+    # 벡터 DB 저장 버튼 및 결과 표시
+    with gr.Row():
+        with gr.Column():
+            save_to_db_button = gr.Button("YOLOv8 인식 결과를 벡터 DB에 저장", variant="primary")
+            save_result = gr.Textbox(label="벡터 DB 저장 결과")
+    
+    # 객체 인식 결과 저장용 상태 변수
+    detection_state = gr.State(None)
+    
     # Set up event handlers
-    yolo_button.click(
+    yolo_result = yolo_button.click(
         fn=process_yolo,
         inputs=input_image,
-        outputs=[yolo_output, yolo_text]
+        outputs=[yolo_output, yolo_text, detection_state]
+    )
+    
+    # 벡터 DB 저장 버튼 이벤트 핸들러
+    save_to_db_button.click(
+        fn=save_objects_to_vector_db,
+        inputs=[input_image, detection_state],
+        outputs=save_result
     )
     
     detr_button.click(