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,44 +10,55 @@ 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, 
-            iou=request.iou_threshold, 
-            conf=request.conf_threshold,
-            classes=request.classes)
+        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))
-    try:
-        response = [{
-            "version": version,
-            "task_type": "det",
-            "shapes": [
-                {
-                    "label": summary['name'],
-                    "color": "#ff0000",
-                    "points": [
-                        [summary['box']['x1'], summary['box']['y1']],
-                        [summary['box']['x2'], summary['box']['y2']]
-                    ],
-                    "group_id": summary['class'],
-                    "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 = []
+    try:
+        for (image, result) in zip(request.image_list, results):
+            label_data:LabelData = {
+                "version": version,
+                "task_type": "det",
+                "shapes": [
+                    {
+                        "label": summary['name'],
+                        "color": "#ff0000",
+                        "points": [
+                            [summary['box']['x1'], summary['box']['y1']],
+                            [summary['box']['x2'], summary['box']['y2']]
+                        ],
+                        "group_id": summary['class'],
+                        "shape_type": "rectangle",
+                        "flags": {}
+                    }    
+                    for summary in result.summary()
+                ],
+                "split": "none",
+                "imageHeight": result.orig_shape[0],
+                "imageWidth": result.orig_shape[1],
+                "imageDepth": 1
+            }
+            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 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)}")