Feat: Detection 오토레이블링 API 구현 - S11P21S002-53

ai/app/api/yolo/detection.py

@@ -0,0 +1,54 @@
+from fastapi import APIRouter, HTTPException
+from schemas.predict_request import PredictRequest
+from schemas.predict_response import PredictResponse
+from services.ai_service import load_detection_model
+from typing import List
+import os
+
+router = APIRouter()
+
+@router.post("/predict", response_model=List[PredictResponse])
+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)
+    try:
+        results = model.predict(
+            source=request.image_path, 
+            iou=request.iou_threshold, 
+            conf=request.conf_threshold,
+            classes=request.classes)
+    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
+        ]
+    except Exception as e:
+        raise HTTPException(status_code=500, detail="label parsing exception: "+str(e))
+    return response

ai/app/main.py

@@ -1,6 +1,12 @@
 from fastapi import FastAPI
-from app.api.endpoints import router
+from api.yolo.detection import router as yolo_detection_router
 
 app = FastAPI()
 
-app.include_router(router)
+# 각 기능별 라우터를 애플리케이션에 등록
+app.include_router(yolo_detection_router, prefix="/api/yolo/detection")
+
+# 애플리케이션 실행
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run("main:app", reload=True)

ai/app/schemas/predict_request.py

@@ -0,0 +1,10 @@
+from pydantic import BaseModel
+from typing import List, Optional
+
+class PredictRequest(BaseModel):
+    projectId: int
+    image_path: str
+    version: Optional[str] = "latest"
+    conf_threshold: Optional[float] = 0.25
+    iou_threshold: Optional[float] = 0.45
+    classes: Optional[List[int]] = None

ai/app/schemas/predict_response.py

@@ -0,0 +1,19 @@
+from pydantic import BaseModel
+from typing import List, Optional, Tuple, Dict
+
+class Shape(BaseModel):
+    label: str
+    color: str
+    points: List[Tuple[float, float]]
+    group_id: Optional[int] = None
+    shape_type: str
+    flags: Dict[str, Optional[bool]] = {}  # key는 문자열, value는 boolean 또는 None
+
+class PredictResponse(BaseModel):
+    version: str
+    task_type: str
+    shapes: List[Shape]
+    split: str
+    imageHeight: int
+    imageWidth: int
+    imageDepth: int

ai/app/services/ai_service.py

@@ -0,0 +1,28 @@
+# ai_service.py
+
+from ultralytics import YOLO  # Ultralytics YOLO 모델을 가져오기
+import os
+
+def load_detection_model(model_path: str = "test/model/initial.pt", device:str ="cpu"): 
+    """
+    지정된 경로에서 YOLO 모델을 로드합니다.
+
+    Args:
+        model_path (str): 모델 파일 경로.
+        device (str): 모델을 로드할 장치. 기본값은 'cpu'.
+                      'cpu' 또는 'cuda'와 같은 장치를 지정할 수 있습니다.
+
+    Returns:
+        YOLO: 로드된 YOLO 모델 인스턴스
+    """
+    
+    if not os.path.exists(model_path):
+        raise FileNotFoundError(f"Model file not found at path: {model_path}")
+
+    try:
+        model = YOLO(model_path)
+        model.to(device)
+        return model
+    except Exception as e:
+        raise RuntimeError(f"Failed to load the model from {model_path}. Error: {str(e)}")
+

ai/requirements.txt

@@ -1,7 +1,7 @@
 fastapi==0.104.1
 uvicorn==0.30.6
-torch==2.4.0
+torch==2.4.0 -f https://download.pytorch.org/whl/cpu
-torchaudio==2.4.0
+torchaudio==2.4.0 -f https://download.pytorch.org/whl/cpu
-torchvision==0.19.0
+torchvision==0.19.0 -f https://download.pytorch.org/whl/cpu
 ultralytics==8.2.82
 ultralytics-thop==2.0.5