diff --git a/ai/.gitignore b/ai/.gitignore index fb43aca..7f82a01 100644 --- a/ai/.gitignore +++ b/ai/.gitignore @@ -30,3 +30,6 @@ dist/ # MacOS 관련 파일 .DS_Store + +# 테스트 파일 +test-data/ \ No newline at end of file diff --git a/ai/app/api/yolo/detection.py b/ai/app/api/yolo/detection.py index 9670a18..8d3e7fd 100644 --- a/ai/app/api/yolo/detection.py +++ b/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)) + + # 추론 결과 -> 레이블 객체 파싱 + response = [] 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 - ] + 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 \ No newline at end of file diff --git a/ai/app/schemas/predict_request.py b/ai/app/schemas/predict_request.py index 023c956..ac2ff3a 100644 --- a/ai/app/schemas/predict_request.py +++ b/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 diff --git a/ai/app/schemas/predict_response.py b/ai/app/schemas/predict_response.py index cd8b6d5..706d896 100644 --- a/ai/app/schemas/predict_response.py +++ b/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 \ No newline at end of file diff --git a/ai/app/services/ai_service.py b/ai/app/services/ai_service.py index 369a1c1..792516b 100644 --- a/ai/app/services/ai_service.py +++ b/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,13 +17,15 @@ 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)}") - + \ No newline at end of file diff --git a/ai/test/images/image_000000001_jpg.rf.02ab6664294833e5f0e89130ecded0b8.jpg b/ai/test/images/image_000000001_jpg.rf.02ab6664294833e5f0e89130ecded0b8.jpg deleted file mode 100644 index 8c0afd3..0000000 Binary files a/ai/test/images/image_000000001_jpg.rf.02ab6664294833e5f0e89130ecded0b8.jpg and /dev/null differ diff --git a/ai/test/images/image_000000002_jpg.rf.8270179e3cd29b97cf502622b381861e.jpg b/ai/test/images/image_000000002_jpg.rf.8270179e3cd29b97cf502622b381861e.jpg deleted file mode 100644 index 237144e..0000000 Binary files a/ai/test/images/image_000000002_jpg.rf.8270179e3cd29b97cf502622b381861e.jpg and /dev/null differ diff --git a/ai/test/images/image_000000003_jpg.rf.db8fd4730b031e35a60e0a60e17a0691.jpg b/ai/test/images/image_000000003_jpg.rf.db8fd4730b031e35a60e0a60e17a0691.jpg deleted file mode 100644 index f4cb350..0000000 Binary files a/ai/test/images/image_000000003_jpg.rf.db8fd4730b031e35a60e0a60e17a0691.jpg and /dev/null differ diff --git a/ai/test/model/initial.pt b/ai/test/model/initial.pt deleted file mode 100644 index 0db4ca4..0000000 Binary files a/ai/test/model/initial.pt and /dev/null differ