Refactor: 요청 응답 객체 구조 변경

ai/.gitignore

@@ -30,3 +30,6 @@ dist/
 
 # MacOS 관련 파일
 .DS_Store
+
+# 테스트 파일
+test-data/

ai/app/api/yolo/detection.py

@@ -1,9 +1,8 @@
 from fastapi import APIRouter, HTTPException
 from schemas.predict_request import PredictRequest
-from schemas.predict_response import PredictResponse
+from schemas.predict_response import PredictResponse, LabelData
 from services.ai_service import load_detection_model
 from typing import List
-import os
 
 router = APIRouter()
 
@@ -11,21 +10,30 @@ router = APIRouter()
 def predict(request: PredictRequest):
     version = "0.1.0"
 
+    # 모델 로드
     try:
         model = load_detection_model()
     except Exception as e:
         raise HTTPException(status_code=500, detail="load model exception: "+str(e))
-    print(model)
+    
+    # 추론
+    results = []
     try:
-        results = model.predict(
-            source=request.image_path, 
+        for image in request.image_list:
+            predict_results = model.predict(
+                source=image.image_url, 
                 iou=request.iou_threshold, 
                 conf=request.conf_threshold,
                 classes=request.classes)
+            results.append(predict_results[0])
     except Exception as e:
         raise HTTPException(status_code=500, detail="model predict exception: "+str(e))
+    
+    # 추론 결과 -> 레이블 객체 파싱
+    response = []
     try:
-        response = [{
+        for (image, result) in zip(request.image_list, results):
+            label_data:LabelData = {
                 "version": version,
                 "task_type": "det",
                 "shapes": [
@@ -40,15 +48,17 @@ def predict(request: PredictRequest):
                         "shape_type": "rectangle",
                         "flags": {}
                     }    
-                
                     for summary in result.summary()
                 ],
                 "split": "none",
                 "imageHeight": result.orig_shape[0],
                 "imageWidth": result.orig_shape[1],
                 "imageDepth": 1
-        } for result in results
-        ]
+            }
+            response.append({
+                "image_id":image.image_id,
+                "data":label_data
+                })
     except Exception as e:
         raise HTTPException(status_code=500, detail="label parsing exception: "+str(e))
     return response

ai/app/schemas/predict_request.py

@@ -1,9 +1,13 @@
 from pydantic import BaseModel
 from typing import List, Optional
 
+class ImageInfo(BaseModel):
+    image_id: int
+    image_url: str    
+
 class PredictRequest(BaseModel):
-    projectId: int
-    image_path: str
+    project_id: int
+    image_list: List[ImageInfo]
     version: Optional[str] = "latest"
     conf_threshold: Optional[float] = 0.25
     iou_threshold: Optional[float] = 0.45

ai/app/schemas/predict_response.py

@@ -7,9 +7,9 @@ class Shape(BaseModel):
     points: List[Tuple[float, float]]
     group_id: Optional[int] = None
     shape_type: str
-    flags: Dict[str, Optional[bool]] = {}  # key는 문자열, value는 boolean 또는 None
+    flags: Dict[str, Optional[bool]] = {}
 
-class PredictResponse(BaseModel):
+class LabelData(BaseModel):
     version: str
     task_type: str
     shapes: List[Shape]
@@ -17,3 +17,7 @@ class PredictResponse(BaseModel):
     imageHeight: int
     imageWidth: int
     imageDepth: int
+
+class PredictResponse(BaseModel):
+    image_id: int
+    data: LabelData

ai/app/services/ai_service.py

@@ -1,9 +1,10 @@
 # ai_service.py
 
 from ultralytics import YOLO # Ultralytics YOLO 모델을 가져오기
+from typing import List
 import os
 
-def load_detection_model(model_path: str = "test/model/initial.pt", device:str ="cpu"): 
+def load_detection_model(model_path: str = "test-data/model/yolov8n.pt", device:str ="cpu"): 
     """
     지정된 경로에서 YOLO 모델을 로드합니다.
 
@@ -16,12 +17,14 @@ def load_detection_model(model_path: str = "test/model/initial.pt", device:str =
         YOLO: 로드된 YOLO 모델 인스턴스
     """
     
-    if not os.path.exists(model_path):
+    if not os.path.exists(model_path) and model_path != "test-data/model/yolov8n.pt":
         raise FileNotFoundError(f"Model file not found at path: {model_path}")
 
     try:
         model = YOLO(model_path)
         model.to(device)
+        # Detection 모델인지 검증
+        # 코드 추가 
         return model
     except Exception as e:
         raise RuntimeError(f"Failed to load the model from {model_path}. Error: {str(e)}")