Merge branch 'ai/refactor/detection' into 'ai/develop'

refactor: detection_train 리팩토링

See merge request s11-s-project/S11P21S002!184

ai/app/api/yolo/detection.py

@@ -8,7 +8,7 @@ from services.create_model import save_model
 from utils.dataset_utils import split_data
 from utils.file_utils import get_dataset_root_path, process_directories, process_image_and_label, join_path
 from utils.slackMessage import send_slack_message
-from utils.api_utils import report_data
+from utils.api_utils import send_data_call_api
 import random
 
 
@@ -59,16 +59,13 @@ def run_predictions(model, image, request, classes):
 # 추론 결과 처리 함수
 def process_prediction_result(result, image, label_map):
     try:
-        random_number = random.randint(0, 0xFFFFFF)
-        color = f"#{random_number:06X}"
-
         label_data = LabelData(
             version="0.0.0",
             task_type="det",
             shapes=[
                 {
                     "label": summary['name'],
-                    "color": color,
+                    "color": get_random_color(),
                     "points": [
                         [summary['box']['x1'], summary['box']['y1']],
                         [summary['box']['x2'], summary['box']['y2']]
@@ -92,6 +89,10 @@ 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")
@@ -103,73 +104,113 @@ async def detection_train(request: TrainRequest, http_request: Request):
     auth_header = http_request.headers.get("Authorization")
     token = auth_header.split(" ")[1] if auth_header and auth_header.startswith("Bearer ") else None
 
+    # 레이블 맵
+    inverted_label_map = {value: key for key, value in request.label_map.items()} if request.label_map else None
+
     # 데이터셋 루트 경로 얻기
     dataset_root_path = get_dataset_root_path(request.project_id)
 
     # 모델 로드
     model = get_model(request)
 
-    # 학습할 모델 카테고리 정리 카테고리가 추가되는 경우에 추가할 수 있게
+    # 학습할 모델 카테고리, 카테고리가 추가되는 경우 추가 작업 필요
-    names = model.names
+    model_categories = model.names
+    
+    # 데이터 전처리
+    preprocess_dataset(dataset_root_path, model_categories, request.data, request.ratio, inverted_label_map)
 
-    # 디렉토리 생성 및 초기화
+    # 학습
-    process_directories(dataset_root_path, names)
+    results = run_train(request,token,model,dataset_root_path)
 
-    # 레이블 맵
+    # last 모델 저장
-    inverted_label_map = {value: key for key, value in request.label_map.items()} if request.label_map else None
+    model_key = save_model(project_id=request.project_id, path=join_path(dataset_root_path, "result", "weights", "best.pt"))
 
-    # 학습 데이터 분류
+    response = {"model_key": model_key, "results": results.results_dict}
-    train_data, val_data = split_data(request.data, request.ratio)
 
+    send_slack_message(f"train 성공{response}", status="success")
+    
+    return response
+
+
+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", inverted_label_map)
+            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", inverted_label_map)
+            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))
+    
+def run_train(request, token, model, dataset_root_path):
+    try:
+        # 데이터 전송 콜백함수
         def send_data(trainer):
-            # 첫번째 epoch는 스킵
+            try:
-            if trainer.epoch == 0:
+                # 첫번째 epoch는 스킵
-                return
+                if trainer.epoch == 0:
+                    return
 
-            ## 남은 시간 계산(초)
+                # 남은 시간 계산(초)
-            left_epochs = trainer.epochs-trainer.epoch
+                left_epochs = trainer.epochs - trainer.epoch
-            left_seconds = left_epochs*trainer.epoch_time
+                left_seconds = left_epochs * trainer.epoch_time
-            ## 로스 box_loss, cls_loss, dfl_loss
-            loss = trainer.label_loss_items(loss_items=trainer.loss_items)
-            data = ReportData(
-                epoch= trainer.epoch,             # 현재 에포크
-                total_epochs= trainer.epochs,     # 전체 에포크
-                box_loss= loss["train/box_loss"], # box loss
-                cls_loss= loss["train/cls_loss"], # cls loss
-                dfl_loss= loss["train/dfl_loss"], # dfl loss
-                fitness= trainer.fitness,         # 적합도
-                epoch_time= trainer.epoch_time,   # 지난 에포크 걸린 시간 (에포크 시작 기준으로 결정)
-                left_seconds= left_seconds        # 남은 시간(초)
-            )
-            # 데이터 전송
-            report_data(request.project_id, request.m_id, data, token)
 
+                # 로스 box_loss, cls_loss, dfl_loss
+                loss = trainer.label_loss_items(loss_items=trainer.loss_items)
+                data = ReportData(
+                    epoch=trainer.epoch,             # 현재 에포크
+                    total_epochs=trainer.epochs,     # 전체 에포크
+                    box_loss=loss["train/box_loss"], # box loss
+                    cls_loss=loss["train/cls_loss"], # cls loss
+                    dfl_loss=loss["train/dfl_loss"], # dfl loss
+                    fitness=trainer.fitness,         # 적합도
+                    epoch_time=trainer.epoch_time,   # 지난 에포크 걸린 시간 (에포크 시작 기준으로 결정)
+                    left_seconds=left_seconds        # 남은 시간(초)
+                )
+                # 데이터 전송
+                send_data_call_api(request.project_id, request.m_id, data, token)
+            except Exception as e:
+                raise HTTPException(status_code=500, detail=f"send_data exception: {e}")
+
+        # 콜백 등록
         model.add_callback("on_train_epoch_start", send_data)
 
-        results = model.train(
+        # 학습 실행
-            data=join_path(dataset_root_path, "dataset.yaml"),
+        try:
-            name=join_path(dataset_root_path, "result"),
+            results = model.train(
-            epochs=request.epochs,
+                data=join_path(dataset_root_path, "dataset.yaml"),
-            batch=request.batch,
+                name=join_path(dataset_root_path, "result"),
-            lr0=request.lr0,
+                epochs=request.epochs,
-            lrf=request.lrf,
+                batch=request.batch,
-            optimizer=request.optimizer
+                lr0=request.lr0,
-        )
+                lrf=request.lrf,
+                optimizer=request.optimizer
+            )
+        except Exception as e:
+            raise HTTPException(status_code=500, detail=f"model train exception: {e}")
 
-        model_key = save_model(project_id=request.project_id, path=join_path(dataset_root_path, "result", "weights", "best.pt"))
+        # 마지막 에포크 전송
-        response = {"model_key": model_key, "results": results.results_dict}
+        model.trainer.epoch += 1
-        send_slack_message(f"train 성공{response}", status="success")
+        send_data(model.trainer)
-        return response
+
+        return results
+
+    except HTTPException as e:
+        raise e # HTTP 예외를 다시 발생
     except Exception as e:
+        raise HTTPException(status_code=500, detail=f"run_train exception: {e}")
 
-        raise HTTPException(status_code=500, detail="model train exception: " + str(e))
 

ai/app/utils/api_utils.py

@@ -3,7 +3,7 @@ from dotenv import load_dotenv
 import os, httpx
 
 
-def report_data(project_id:int, model_id:int, data:ReportData, token):
+def send_data_call_api(project_id:int, model_id:int, data:ReportData, token):
     try:
         load_dotenv() 
         # main.py와 같은 디렉토리에 .env 파일 생성해서 따옴표 없이 입력