From 27cfd01fd7b1cf40540b45cc4f159375c3f00d0c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=EA=B9=80=EC=A7=84=ED=98=84?= Date: Wed, 25 Sep 2024 15:24:14 +0900 Subject: [PATCH] =?UTF-8?q?Fix:=20=EB=A6=AC=ED=8C=A9=ED=86=A0=EB=A7=81?= =?UTF-8?q?=EC=97=90=EC=84=9C=20=EC=83=9D=EA=B8=B4=20=EB=B2=84=EA=B7=B8=20?= =?UTF-8?q?=EC=88=98=EC=A0=95?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- ai/app/api/yolo/detection.py | 15 ++++++++------- ai/app/api/yolo/segmentation.py | 2 +- ai/app/main.py | 8 ++++---- ai/app/schemas/train_request.py | 6 +++--- ai/app/services/load_model.py | 4 +--- ai/app/utils/file_utils.py | 21 ++++++++++----------- 6 files changed, 27 insertions(+), 29 deletions(-) diff --git a/ai/app/api/yolo/detection.py b/ai/app/api/yolo/detection.py index 845fc44..ff542d3 100644 --- a/ai/app/api/yolo/detection.py +++ b/ai/app/api/yolo/detection.py @@ -18,7 +18,7 @@ router = APIRouter() @router.post("/predict") async def detection_predict(request: PredictRequest): - send_slack_message(f"predict 요청{request}", status="success") + # send_slack_message(f"predict 요청{request}", status="success") # 모델 로드 model = get_model(request) @@ -33,6 +33,8 @@ async def detection_predict(request: PredictRequest): # 추론 try: results = run_predictions(model, url_list, request, classes) + print(len(results)) + print(len(request.image_list)) response = [process_prediction_result(result, image, request.label_map) for result, image in zip(results,request.image_list)] return response @@ -51,13 +53,12 @@ def get_model(request: PredictRequest): # 추론 실행 함수 def run_predictions(model, image, request, classes): try: - predict_results = model.predict( - source=image, - iou=request.iou_threshold, - conf=request.conf_threshold, - classes=classes + return model.predict( + source=image, + iou=request.iou_threshold, + conf=request.conf_threshold, + classes=classes ) - return predict_results[0] except Exception as e: raise HTTPException(status_code=500, detail="model predict exception: " + str(e)) diff --git a/ai/app/api/yolo/segmentation.py b/ai/app/api/yolo/segmentation.py index 686e451..b0ca826 100644 --- a/ai/app/api/yolo/segmentation.py +++ b/ai/app/api/yolo/segmentation.py @@ -11,7 +11,7 @@ def predict(request: PredictRequest): # 모델 로드 try: - model = load_segmentation_model() + model = load_segmentation_model(request.project_id, request.m_key) except Exception as e: raise HTTPException(status_code=500, detail="load model exception: "+str(e)) diff --git a/ai/app/main.py b/ai/app/main.py index a346b5e..a30ccea 100644 --- a/ai/app/main.py +++ b/ai/app/main.py @@ -10,7 +10,7 @@ app.include_router(yolo_detection_router, prefix="/api/detection", tags=["Detect app.include_router(yolo_segmentation_router, prefix="/api/segmentation", tags=["Segmentation"]) app.include_router(yolo_model_router, prefix="/api/model", tags=["Model"]) -# 애플리케이션 실행 -if __name__ == "__main__": - import uvicorn - uvicorn.run("main:app", reload=True) +# # 애플리케이션 실행 +# if __name__ == "__main__": +# import uvicorn +# uvicorn.run("main:app", reload=True) diff --git a/ai/app/schemas/train_request.py b/ai/app/schemas/train_request.py index 934d8c8..f60a538 100644 --- a/ai/app/schemas/train_request.py +++ b/ai/app/schemas/train_request.py @@ -8,9 +8,9 @@ class TrainDataInfo(BaseModel): class TrainRequest(BaseModel): project_id: int - m_key: Optional[str] = Field(None, alias="model_key") - m_id: Optional[str] = Field(None, alias="model_id") # 학습 중 에포크 결과를 보낼때 model_id를 보냄 - label_map: dict[int, int] = Field(None, description="모델 레이블 카테고리 idx: 프로젝트 레이블 카테고리 idx , None 일경우 레이블 데이터(프로젝트 레이블)의 idx로 학습") + m_key: str = Field("yolo8", alias="model_key") + m_id: int = Field(..., alias="model_id") # 학습 중 에포크 결과를 보낼때 model_id를 보냄 + label_map: dict[int, int] = Field({}, description="모델 레이블 카테고리 idx: 프로젝트 레이블 카테고리 idx , None 일경우 레이블 데이터(프로젝트 레이블)의 idx로 학습") data: List[TrainDataInfo] ratio: float = 0.8 # 훈련/검증 분할 비율 diff --git a/ai/app/services/load_model.py b/ai/app/services/load_model.py index 0c60b50..e799353 100644 --- a/ai/app/services/load_model.py +++ b/ai/app/services/load_model.py @@ -1,8 +1,6 @@ # ai_service.py from ultralytics import YOLO # Ultralytics YOLO 모델을 가져오기 -from ultralytics.models.yolo.model import YOLO as YOLO_Model -from ultralytics.nn.tasks import DetectionModel, SegmentationModel import os import torch @@ -12,7 +10,7 @@ def load_detection_model(project_id:int, model_key:str): if model_key in default_model_map: model = YOLO(default_model_map[model_key]) else: - model = load_model(model_path=os.path.join("projects",str(project_id),"models",model_key)) + model = load_model(model_path=os.path.join("projects",str(project_id),"models", model_key)) # Detection 모델인지 검증 if model.task != "detect": diff --git a/ai/app/utils/file_utils.py b/ai/app/utils/file_utils.py index 81ed9e3..a68a2d5 100644 --- a/ai/app/utils/file_utils.py +++ b/ai/app/utils/file_utils.py @@ -40,7 +40,6 @@ def process_directories(dataset_root_path:str, names:list[str]): make_yml(dataset_root_path, names) def process_image_and_label(data:TrainDataInfo, dataset_root_path:str, child_path:str, label_map:dict[int, int]|None): - """이미지 저장 및 레이블 파일 생성""" # 이미지 url로부터 파일명 분리 img_name = data.image_url.split('/')[-1] @@ -64,17 +63,17 @@ def process_image_and_label(data:TrainDataInfo, dataset_root_path:str, child_pat def create_detection_train_label(label:LabelData, label_path:str, label_map:dict[int, int]|None): with open(label_path, "w") as train_label_txt: - for shape in label.shapes: + for shape in label["shapes"]: train_label = [] - x1 = shape.points[0][0] - y1 = shape.points[0][1] - x2 = shape.points[1][0] - y2 = shape.points[1][1] - train_label.append(str(label_map[shape.group_id]) if label_map else str(shape.group_id)) # label Id - 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 )) # 높이 + x1 = shape["points"][0][0] + y1 = shape["points"][0][1] + x2 = shape["points"][1][0] + y2 = shape["points"][1][1] + train_label.append(str(label_map[shape["group_id"]]) if label_map else str(shape["group_id"])) # label Id + 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):