From 75a12dd10c0736fbe0c102013c07f0047b0bc921 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=EA=B9=80=EC=A7=84=ED=98=84?= Date: Mon, 9 Sep 2024 17:46:15 +0900 Subject: [PATCH] =?UTF-8?q?Feat:=20Detection=20=EB=AA=A8=EB=8D=B8=20?= =?UTF-8?q?=ED=95=99=EC=8A=B5=20API=20=EA=B5=AC=ED=98=84=20-=20S11P21S002-?= =?UTF-8?q?117?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- ai/.gitignore | 7 ++++++- ai/app/api/yolo/detection.py | 17 ++++++++++++++--- ai/app/schemas/train_request.py | 6 ++++-- ai/app/services/ai_service.py | 2 +- ai/app/utils/file_utils.py | 26 ++++++++++++++++---------- 5 files changed, 41 insertions(+), 17 deletions(-) diff --git a/ai/.gitignore b/ai/.gitignore index 7f82a01..7c7931f 100644 --- a/ai/.gitignore +++ b/ai/.gitignore @@ -32,4 +32,9 @@ dist/ .DS_Store # 테스트 파일 -test-data/ \ No newline at end of file +test-data/ + +# 리소스 +resources/ +datasets/ +*.pt \ No newline at end of file diff --git a/ai/app/api/yolo/detection.py b/ai/app/api/yolo/detection.py index 07b33b6..82ac1ac 100644 --- a/ai/app/api/yolo/detection.py +++ b/ai/app/api/yolo/detection.py @@ -4,9 +4,9 @@ from schemas.train_request import TrainRequest from schemas.predict_response import PredictResponse, LabelData from services.ai_service import load_detection_model from utils.dataset_utils import split_data -from utils.file_utils import get_dataset_root_path, process_directories, process_image_and_label +from utils.file_utils import get_dataset_root_path, process_directories, process_image_and_label, join_path from typing import List -from PIL import Image +from fastapi.responses import FileResponse router = APIRouter() @router.post("/detection", response_model=List[PredictResponse]) @@ -96,4 +96,15 @@ def train(request: TrainRequest): # 검증 데이터 처리 for data in val_data: - process_image_and_label(data, dataset_root_path, "val") \ No newline at end of file + process_image_and_label(data, dataset_root_path, "val") + + model = load_detection_model("test-data/model/best.pt") + + model.train( + data=join_path(dataset_root_path,"dataset.yaml"), + name=join_path(dataset_root_path,"result"), + epochs= request.epochs, + batch=request.batch, + ) + + return FileResponse(path=join_path(dataset_root_path, "result", "weights", "best.pt"), filename="best.pt", media_type="application/octet-stream") \ No newline at end of file diff --git a/ai/app/schemas/train_request.py b/ai/app/schemas/train_request.py index a239dd2..392ca49 100644 --- a/ai/app/schemas/train_request.py +++ b/ai/app/schemas/train_request.py @@ -1,5 +1,5 @@ from pydantic import BaseModel -from typing import List, Optional +from typing import List, Optional, Union from schemas.predict_response import LabelData class TrainDataInfo(BaseModel): @@ -11,4 +11,6 @@ class TrainRequest(BaseModel): project_id: int data: List[TrainDataInfo] seed: Optional[int] = None # 랜덤 변수 시드 - ratio: Optional[float] = 0.8 # 훈련/검증 분할 비율 + ratio: float = 0.8 # 훈련/검증 분할 비율 + epochs: int = 50 # 훈련 반복 횟수 + batch: Union[float, int] = -1 # 훈련 batch 수[int] or GPU의 사용률 자동[float] default(-1): gpu의 60% 사용 유지 diff --git a/ai/app/services/ai_service.py b/ai/app/services/ai_service.py index 790f61d..234948b 100644 --- a/ai/app/services/ai_service.py +++ b/ai/app/services/ai_service.py @@ -6,7 +6,7 @@ from ultralytics.nn.tasks import DetectionModel, SegmentationModel import os import torch -def load_detection_model(model_path: str = "test-data/model/yolov8n.pt", device:str ="auto"): +def load_detection_model(model_path: str = os.path.join("test-data","model","initial.pt"), device:str ="auto"): """ 지정된 경로에서 YOLO 모델을 로드합니다. diff --git a/ai/app/utils/file_utils.py b/ai/app/utils/file_utils.py index f9605ca..ffb32fb 100644 --- a/ai/app/utils/file_utils.py +++ b/ai/app/utils/file_utils.py @@ -5,8 +5,8 @@ from PIL import Image from schemas.train_request import TrainDataInfo def get_dataset_root_path(project_id): - """프로젝트 ID를 기반으로 데이터셋 루트 경로 반환""" - return os.path.join('test-data', 'projects', str(project_id), 'train_model') + """데이터셋 루트 절대 경로 반환""" + return os.path.join(os.getcwd(), 'datasets', 'train') def make_dir(path:str, init: bool): """ @@ -19,8 +19,8 @@ def make_dir(path:str, init: bool): def make_yml(path:str): data = { - "train": "train", - "val": "val", + "train": f"{path}/train", + "val": f"{path}/val", "nc": 80, "names": { @@ -106,7 +106,7 @@ def make_yml(path:str): 79: "toothbrush" } } - with open(path, 'w') as f: + with open(os.path.join(path, "dataset.yaml"), 'w') as f: yaml.dump(data, f) def process_directories(dataset_root_path:str): @@ -114,7 +114,9 @@ def process_directories(dataset_root_path:str): make_dir(dataset_root_path, init=False) make_dir(os.path.join(dataset_root_path, "train"), init=True) make_dir(os.path.join(dataset_root_path, "val"), init=True) - make_yml(os.path.join(dataset_root_path, "dataset.yaml")) + if os.path.exists(os.path.join(dataset_root_path, "result")): + shutil.rmtree(os.path.join(dataset_root_path, "result")) + make_yml(dataset_root_path) def process_image_and_label(data:TrainDataInfo, dataset_root_path:str, child_path:str): @@ -138,8 +140,12 @@ def process_image_and_label(data:TrainDataInfo, dataset_root_path:str, child_pat x2 = shape.points[1][0] y2 = shape.points[1][1] train_label.append(str(shape.group_id)) # label Id - train_label.append(str((x1 + x2) / 2)) # 중심 x 좌표 - train_label.append(str((y1 + y2) / 2)) # 중심 y 좌표 - train_label.append(str(x2 - x1)) # 너비 - train_label.append(str(y2 - y1)) # 높이 + train_label.append(str((x1 + x2) / 2 / label.imageWidth)) # 중심 x 좌표 + train_label.append(str((y1 + y2) / 2 / label.imageHeight)) # 중심 y 좌표 + train_label.append(str((x2 - x1) / label.imageWidth)) # 너비 + train_label.append(str((y2 - y1) / label.imageHeight )) # 높이 train_label_txt.write(" ".join(train_label)+"\n") + +def join_path(path, *paths): + """os.path.join()과 같은 기능, os import 하기 싫어서 만듦""" + return os.path.join(path, *paths) \ No newline at end of file