diff --git a/ai/app/api/yolo/detection.py b/ai/app/api/yolo/detection.py index eb3f974..1a659de 100644 --- a/ai/app/api/yolo/detection.py +++ b/ai/app/api/yolo/detection.py @@ -116,7 +116,7 @@ async def detection_train(request: TrainRequest): # 이 값을 학습할때 넣으면 이 카테고리들이 학습됨 names = list(request.label_map) - # 데이터 전처리: 학습할 디렉토리 & 데이터셋 를 생성 + # 데이터 전처리: 학습할 디렉토리 & 데이터셋 설정 파일을 생성 process_directories(dataset_root_path, names) # 데이터 전처리: 데이터를 학습데이터와 검증데이터로 분류 @@ -183,6 +183,7 @@ def run_train(request, model, dataset_root_path): data = ReportData( epoch=trainer.epoch, # 현재 에포크 total_epochs=trainer.epochs, # 전체 에포크 + seg_loss=0, # seg_loss box_loss=loss["train/box_loss"], # box loss cls_loss=loss["train/cls_loss"], # cls loss dfl_loss=loss["train/dfl_loss"], # dfl loss diff --git a/ai/app/api/yolo/segmentation.py b/ai/app/api/yolo/segmentation.py index 2c47e22..79d4fff 100644 --- a/ai/app/api/yolo/segmentation.py +++ b/ai/app/api/yolo/segmentation.py @@ -1,4 +1,5 @@ -from fastapi import APIRouter, HTTPException, Request +from fastapi import APIRouter, HTTPException +from api.yolo.detection import get_classes, run_predictions, get_random_color, split_data, download_data from schemas.predict_request import PredictRequest from schemas.train_request import TrainRequest from schemas.predict_response import PredictResponse, LabelData @@ -15,25 +16,24 @@ router = APIRouter() @router.post("/predict") async def segmentation_predict(request: PredictRequest): - - send_slack_message(f"seg predict 요청: {request}", status="success") + send_slack_message(f"predict 요청: {request}", status="success") # 모델 로드 - model = get_model(request) - - # 모델 레이블 카테고리 연결 - classes = list(request.label_map) if request.label_map else None + model = get_model(request.project_id, request.m_key) # 이미지 데이터 정리 url_list = list(map(lambda x:x.image_url, request.image_list)) + # 이 값을 모델에 입력하면 해당하는 클래스 id만 출력됨 + classes = get_classes(request.label_map, model.names) + # 추론 results = run_predictions(model, url_list, request, classes) # 추론 결과 변환 response = [process_prediction_result(result, image, request.label_map) for result, image in zip(results,request.image_list)] send_slack_message(f"predict 성공{response}", status="success") - return response + return response # 모델 로드 def get_model(request: PredictRequest): @@ -42,19 +42,6 @@ def get_model(request: PredictRequest): except Exception as e: raise HTTPException(status_code=500, detail="load model exception: " + str(e)) -# 추론 실행 함수 -def run_predictions(model, image, request, classes): - try: - return model.predict( - source=image, - iou=request.iou_threshold, - conf=request.conf_threshold, - classes=classes - ) - except Exception as e: - raise HTTPException(status_code=500, detail="model predict exception: " + str(e)) - - # 추론 결과 처리 함수 def process_prediction_result(result, image, label_map): try: @@ -66,7 +53,7 @@ def process_prediction_result(result, image, label_map): "label": summary['name'], "color": get_random_color(), "points": list(zip(summary['segments']['x'], summary['segments']['y'])), - "group_id": label_map[summary['class']] if label_map else summary['class'], + "group_id": label_map[summary['name']], "shape_type": "polygon", "flags": {} } @@ -85,80 +72,49 @@ def process_prediction_result(result, image, label_map): data=label_data.model_dump_json() ) -def get_random_color(): - random_number = random.randint(0, 0xFFFFFF) - return f"#{random_number:06X}" - @router.post("/train") async def segmentation_train(request: TrainRequest): - + send_slack_message(f"train 요청{request}", status="success") - try: - # 레이블 맵 - inverted_label_map = {value: key for key, value in request.label_map.items()} if request.label_map else None + # 데이터셋 루트 경로 얻기 (프로젝트 id 기반) + dataset_root_path = get_dataset_root_path(request.project_id) - # 데이터셋 루트 경로 얻기 - dataset_root_path = get_dataset_root_path(request.project_id) + # 모델 로드 + model = get_model(request.project_id, request.m_key) - # 모델 로드 - model = get_model(request) + # 이 값을 학습할때 넣으면 이 카테고리들이 학습됨 + names = list(request.label_map) + + # 데이터 전처리: 학습할 디렉토리 & 데이터셋 설정 파일을 생성 + process_directories(dataset_root_path, names) - # 학습할 모델 카테고리, 카테고리가 추가되는 경우 추가 작업 필요 - model_categories = model.names - - # 데이터 전처리 - preprocess_dataset(dataset_root_path, model_categories, request.data, request.ratio, inverted_label_map) + # 데이터 전처리: 데이터를 학습데이터와 검증데이터로 분류 + train_data, val_data = split_data(request.data, request.ratio) - # 학습 - results = run_train(request, model,dataset_root_path) + # 데이터 전처리: 데이터 이미지 및 레이블 다운로드 + download_data(train_data, val_data, dataset_root_path) - # best 모델 저장 - model_key = save_model(project_id=request.project_id, path=join_path(dataset_root_path, "result", "weights", "best.pt")) - - result = results.results_dict + # 학습 + results = run_train(request, model,dataset_root_path) - response = TrainResponse( - modelKey=model_key, - precision= result["metrics/precision(M)"], - recall= result["metrics/recall(M)"], - mAP50= result["metrics/mAP50(M)"], - mAP5095= result["metrics/mAP50-95(M)"], - fitness= result["fitness"] - ) - send_slack_message(f"train 성공{response}", status="success") + # best 모델 저장 + model_key = save_model(project_id=request.project_id, path=join_path(dataset_root_path, "result", "weights", "best.pt")) + + result = results.results_dict + + response = TrainResponse( + modelKey=model_key, + precision= result["metrics/precision(M)"], + recall= result["metrics/recall(M)"], + mAP50= result["metrics/mAP50(M)"], + mAP5095= result["metrics/mAP50-95(M)"], + fitness= result["fitness"] + ) + send_slack_message(f"train 성공{response}", status="success") - return response - - except HTTPException as e: - raise e - except Exception as e: - raise HTTPException(status_code=500, detail=str(e)) - - -def preprocess_dataset(dataset_root_path, model_categories, data, ratio, label_map): - try: - # 디렉토리 생성 및 초기화 - process_directories(dataset_root_path, model_categories) - - # 학습 데이터 분류 - train_data, val_data = split_data(data, ratio) - if not train_data or not val_data: - raise HTTPException(status_code=400, detail="data split exception: data size is too small or \"ratio\" has invalid value") - - # 학습 데이터 처리 - for data in train_data: - process_image_and_label(data, dataset_root_path, "train", label_map) - - # 검증 데이터 처리 - for data in val_data: - process_image_and_label(data, dataset_root_path, "val", label_map) - - except HTTPException as e: - raise e # HTTP 예외를 다시 발생 - except Exception as e: - raise HTTPException(status_code=500, detail="preprocess dataset exception: " + str(e)) + return response def run_train(request, model, dataset_root_path): try: @@ -178,9 +134,10 @@ def run_train(request, model, dataset_root_path): data = ReportData( epoch=trainer.epoch, # 현재 에포크 total_epochs=trainer.epochs, # 전체 에포크 - box_loss=loss["train/box_loss"], # box loss + seg_loss=loss["train/seg_loss"], # seg_loss + box_loss=0, # box loss cls_loss=loss["train/cls_loss"], # cls loss - dfl_loss=loss["train/dfl_loss"], # dfl loss + dfl_loss=0, # dfl loss fitness=trainer.fitness, # 적합도 epoch_time=trainer.epoch_time, # 지난 에포크 걸린 시간 (에포크 시작 기준으로 결정) left_seconds=left_seconds # 남은 시간(초) diff --git a/ai/app/schemas/train_report_data.py b/ai/app/schemas/train_report_data.py index 600c129..b4bfdc0 100644 --- a/ai/app/schemas/train_report_data.py +++ b/ai/app/schemas/train_report_data.py @@ -4,6 +4,7 @@ class ReportData(BaseModel): epoch: int # 현재 에포크 total_epochs: int # 전체 에포크 + seg_loss: float # seg_loss box_loss: float # box loss cls_loss: float # cls loss dfl_loss: float # dfl loss