diff --git a/ai/app/api/yolo/model.py b/ai/app/api/yolo/model.py index 489912f..f166819 100644 --- a/ai/app/api/yolo/model.py +++ b/ai/app/api/yolo/model.py @@ -32,9 +32,6 @@ def get_model_list(project_id:int): @router.post("/projects/{project_id}", status_code=201) def create_model(project_id: int, request: ModelCreateRequest): - if request.project_type not in ["segmentation", "detection", "classification"]: - raise HTTPException(status_code=400, - detail= f"Invalid type '{request.type}'. Must be one of \"segmentation\", \"detection\", \"classification\".") model_key = create_new_model(project_id, request.project_type, request.pretrained) return {"model_key": model_key} diff --git a/ai/app/schemas/model_create_request.py b/ai/app/schemas/model_create_request.py index 70aa483..0f17676 100644 --- a/ai/app/schemas/model_create_request.py +++ b/ai/app/schemas/model_create_request.py @@ -1,5 +1,6 @@ from pydantic import BaseModel +from typing import Literal class ModelCreateRequest(BaseModel): - project_type: str + project_type: Literal["segmentation", "detection", "classification"] pretrained:bool = True \ No newline at end of file diff --git a/ai/app/schemas/predict_request.py b/ai/app/schemas/predict_request.py index ba02092..e4cf3b3 100644 --- a/ai/app/schemas/predict_request.py +++ b/ai/app/schemas/predict_request.py @@ -1,5 +1,4 @@ from pydantic import BaseModel, Field -from typing import Optional, Union class ImageInfo(BaseModel): image_id: int @@ -9,7 +8,7 @@ class ImageInfo(BaseModel): class PredictRequest(BaseModel): project_id: int m_key: str = Field("yolo8", alias="model_key") # model_ 로 시작하는 변수를 BaseModel의 변수로 만들경우 Warning 떠서 m_key로 대체 - label_map: dict[str, int] = Field(..., description="프로젝트 레이블 이름: 프로젝트 레이블 pk , None일 경우 모델 레이블 카테고리 idx로 레이블링") + label_map: dict[str, int] = Field(..., description="프로젝트 레이블 이름: 프로젝트 레이블 pk") image_list: list[ImageInfo] # 이미지 리스트 - conf_threshold: float = 0.25 # - iou_threshold: float = 0.45 + conf_threshold: float = Field(0.25, gt=0, lt= 1) + iou_threshold: float = Field(0.45, gt=0, lt= 1) diff --git a/ai/app/schemas/train_report_data.py b/ai/app/schemas/train_report_data.py index b4bfdc0..1e4a89f 100644 --- a/ai/app/schemas/train_report_data.py +++ b/ai/app/schemas/train_report_data.py @@ -1,7 +1,6 @@ from pydantic import BaseModel class ReportData(BaseModel): - epoch: int # 현재 에포크 total_epochs: int # 전체 에포크 seg_loss: float # seg_loss diff --git a/ai/app/schemas/train_request.py b/ai/app/schemas/train_request.py index f090219..a06899f 100644 --- a/ai/app/schemas/train_request.py +++ b/ai/app/schemas/train_request.py @@ -1,22 +1,22 @@ from pydantic import BaseModel, Field -from typing import List, Optional, Union, Literal -from schemas.predict_response import LabelData +from typing import Literal class TrainDataInfo(BaseModel): image_url: str data_url: str class TrainRequest(BaseModel): - project_id: int + project_id: int = Field(..., gt= 0) m_key: str = Field("yolo8", alias="model_key") - m_id: int = Field(..., alias="model_id") # 학습 중 에포크 결과를 보낼때 model_id를 보냄 - label_map: dict[str, int] = Field(..., description="프로젝트 레이블 이름: 프로젝트 레이블 pk , None일 경우 모델 레이블 카테고리 idx로 레이블링") - data: List[TrainDataInfo] - ratio: float = 0.8 # 훈련/검증 분할 비율 + m_id: int = Field(..., alias="model_id", gt= 0) # 학습 중 에포크 결과를 보낼때 model_id를 보냄 + label_map: dict[str, int] = Field(..., description="프로젝트 레이블 이름: 프로젝트 레이블 pk") + data: list[TrainDataInfo] + ratio: float = Field(0.8, gt=0, lt=1) # 훈련/검증 분할 비율 # 학습 파라미터 - epochs: int = 50 # 훈련 반복 횟수 - batch: Union[float, int] = -1 # 훈련 batch 수[int] or GPU의 사용률 자동[float] default(-1): gpu의 60% 사용 유지 - lr0: float = 0.01 # 초기 학습 가중치 - lrf: float = 0.01 # lr0 기준으로 학습 가중치의 최종 수렴치 (ex lr0의 0.01배) + epochs: int = Field(50, gt= 0, lt = 1000) # 훈련 반복 횟수 + batch: int = Field(16, gt=0, le = 10000) # 훈련 batch 수[int] or GPU의 사용률 자동[float] default(-1): gpu의 60% 사용 유지 + lr0: float = Field(0.01, gt= 0, lt= 1) # 초기 학습 가중치 + lrf: float = Field(0.01, gt= 0, lt= 1) # lr0 기준으로 학습 가중치의 최종 수렴치 (ex lr0의 0.01배) optimizer: Literal['auto', 'SGD', 'Adam', 'AdamW', 'NAdam', 'RAdam', 'RMSProp'] = 'auto' + diff --git a/ai/app/utils/file_utils.py b/ai/app/utils/file_utils.py index 15faa2d..7cab7df 100644 --- a/ai/app/utils/file_utils.py +++ b/ai/app/utils/file_utils.py @@ -141,6 +141,10 @@ def process_image_and_label_in_cls(data:TrainDataInfo, dataset_root_path:str, ch # 레이블 객체 불러오기 label = json.loads(urllib.request.urlopen(data.data_url).read()) + if not label["shapes"]: + # assert label["shapes"], No Label. Failed Download" # AssertionError 발생 + print("No Label. Failed Download") + return label_name = label["shapes"][0]["label"] label_path = os.path.join(dataset_root_path,child_path,label_name) @@ -149,8 +153,8 @@ def process_image_and_label_in_cls(data:TrainDataInfo, dataset_root_path:str, ch if os.path.exists(label_path): urllib.request.urlretrieve(data.image_url, os.path.join(label_path, img_name)) else: - # raise FileNotFoundError("failed download") - print("Not Found Label Category. Failed Download") + # raise FileNotFoundError("No Label Category. Failed Download") + print("No Label Category. Failed Download") # 레이블 데이터 중에서 프로젝트 카테고리에 해당되지않는 데이터가 있는 경우 처리 1. 에러 raise 2. 무시(+ warning) \ No newline at end of file