From 4b6751a00bb23db4d657ae3d98787459308d9afb Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=EA=B9=80=EC=A7=84=ED=98=84?= <kimjh9903@naver.com>
Date: Fri, 27 Sep 2024 14:18:59 +0900
Subject: [PATCH] =?UTF-8?q?Refactor:=20Train=20response=EC=97=90=20accurac?=
 =?UTF-8?q?y=20=EC=B6=94=EA=B0=80?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 ai/app/api/yolo/classfication.py | 11 ++++++++---
 ai/app/api/yolo/detection.py     |  1 +
 ai/app/api/yolo/segmentation.py  |  1 +
 ai/app/schemas/train_response.py |  4 ----
 4 files changed, 10 insertions(+), 7 deletions(-)

diff --git a/ai/app/api/yolo/classfication.py b/ai/app/api/yolo/classfication.py
index a258cf4..8496d14 100644
--- a/ai/app/api/yolo/classfication.py
+++ b/ai/app/api/yolo/classfication.py
@@ -4,7 +4,7 @@ from schemas.predict_request import PredictRequest
 from schemas.train_request import TrainRequest, TrainDataInfo
 from schemas.predict_response import PredictResponse, LabelData
 from schemas.train_report_data import ReportData
-from schemas.train_response import ClassificationTrainResponse
+from schemas.train_response import TrainResponse
 from services.load_model import load_classification_model
 from services.create_model import save_model
 from utils.file_utils import get_dataset_root_path, process_directories_in_cls, process_image_and_label_in_cls, join_path
@@ -105,11 +105,16 @@ async def classification_train(request: TrainRequest):
     
     result = results.results_dict
 
-    response = ClassificationTrainResponse(
+    response = TrainResponse(
         modelKey=model_key,
-        precision= result["accuracy_top1"],
+        precision= 0,
+        recall= 0,
+        mAP50= 0,
+        mAP5095= 0,
+        accuracy=result["accuracy_top1"],
         fitness= result["fitness"]
     )
+    
     send_slack_message(f"train 성공{response}", status="success")
             
     return response
diff --git a/ai/app/api/yolo/detection.py b/ai/app/api/yolo/detection.py
index c30282a..b3ca8f7 100644
--- a/ai/app/api/yolo/detection.py
+++ b/ai/app/api/yolo/detection.py
@@ -144,6 +144,7 @@ async def detection_train(request: TrainRequest):
         recall= result["metrics/recall(B)"],
         mAP50= result["metrics/mAP50(B)"],
         mAP5095= result["metrics/mAP50-95(B)"],
+        accuracy=0,
         fitness= result["fitness"]
     )
     send_slack_message(f"train 성공{response}", status="success")
diff --git a/ai/app/api/yolo/segmentation.py b/ai/app/api/yolo/segmentation.py
index 5df4ca1..89458ed 100644
--- a/ai/app/api/yolo/segmentation.py
+++ b/ai/app/api/yolo/segmentation.py
@@ -114,6 +114,7 @@ async def segmentation_train(request: TrainRequest):
         recall= result["metrics/recall(M)"],
         mAP50= result["metrics/mAP50(M)"],
         mAP5095= result["metrics/mAP50-95(M)"],
+        accuracy = 0,
         fitness= result["fitness"]
     )
     send_slack_message(f"train 성공{response}", status="success")
diff --git a/ai/app/schemas/train_response.py b/ai/app/schemas/train_response.py
index 3a64a4b..222b3ce 100644
--- a/ai/app/schemas/train_response.py
+++ b/ai/app/schemas/train_response.py
@@ -6,9 +6,5 @@ class TrainResponse(BaseModel):
     recall: float
     mAP50: float
     mAP5095: float
-    fitness: float
-
-class ClassificationTrainResponse(BaseModel):
-    modelKey: str
     accuracy: float
     fitness: float
\ No newline at end of file