From 02929ebc86ffbb377c326783d55994333b48aa0e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=EA=B9=80=EC=A7=84=ED=98=84?= Date: Fri, 27 Sep 2024 12:58:18 +0900 Subject: [PATCH] =?UTF-8?q?Fix:=20Detection,=20Segmentation=20=ED=95=99?= =?UTF-8?q?=EC=8A=B5=20=EB=A0=88=EC=9D=B4=EB=B8=94=20=EC=BB=A8=EB=B2=84?= =?UTF-8?q?=ED=84=B0=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/detection.py | 13 +++++++++---- ai/app/api/yolo/segmentation.py | 10 +++++++--- ai/app/utils/file_utils.py | 19 ++++++++----------- 3 files changed, 24 insertions(+), 18 deletions(-) diff --git a/ai/app/api/yolo/detection.py b/ai/app/api/yolo/detection.py index 1a659de..c30282a 100644 --- a/ai/app/api/yolo/detection.py +++ b/ai/app/api/yolo/detection.py @@ -115,6 +115,11 @@ async def detection_train(request: TrainRequest): # 이 값을 학습할때 넣으면 이 카테고리들이 학습됨 names = list(request.label_map) + + # 레이블 변환기 (file_util.py/create_detection_train_label() 에 쓰임) + label_converter = {request.label_map[key]:idx for idx, key in enumerate(request.label_map)} + # key : 데이터에 저장된 프로젝트 카테고리 id + # value : 모델에 저장될 카테고리 id (모델에는 key의 idx 순서대로 저장될 것임) # 데이터 전처리: 학습할 디렉토리 & 데이터셋 설정 파일을 생성 process_directories(dataset_root_path, names) @@ -123,7 +128,7 @@ async def detection_train(request: TrainRequest): train_data, val_data = split_data(request.data, request.ratio) # 데이터 전처리: 데이터 이미지 및 레이블 다운로드 - download_data(train_data, val_data, dataset_root_path) + download_data(train_data, val_data, dataset_root_path, label_converter) # 학습 results = run_train(request, model,dataset_root_path) @@ -155,13 +160,13 @@ def split_data(data:list[TrainDataInfo], ratio:float): except Exception as e: raise HTTPException(status_code=500, detail="exception in split_data(): " + str(e)) -def download_data(train_data:list[TrainDataInfo], val_data:list[TrainDataInfo], dataset_root_path:str): +def download_data(train_data:list[TrainDataInfo], val_data:list[TrainDataInfo], dataset_root_path:str, label_converter:dict[int, int]): try: for data in train_data: - process_image_and_label(data, dataset_root_path, "train") + process_image_and_label(data, dataset_root_path, "train", label_converter) for data in val_data: - process_image_and_label(data, dataset_root_path, "val") + process_image_and_label(data, dataset_root_path, "val", label_converter) except Exception as e: raise HTTPException(status_code=500, detail="exception in download_data(): " + str(e)) diff --git a/ai/app/api/yolo/segmentation.py b/ai/app/api/yolo/segmentation.py index 79d4fff..5df4ca1 100644 --- a/ai/app/api/yolo/segmentation.py +++ b/ai/app/api/yolo/segmentation.py @@ -7,10 +7,9 @@ from schemas.train_report_data import ReportData from schemas.train_response import TrainResponse from services.load_model import load_segmentation_model from services.create_model import save_model -from utils.file_utils import get_dataset_root_path, process_directories, process_image_and_label, join_path +from utils.file_utils import get_dataset_root_path, process_directories, join_path from utils.slackMessage import send_slack_message from utils.api_utils import send_data_call_api -import random router = APIRouter() @@ -86,6 +85,11 @@ async def segmentation_train(request: TrainRequest): # 이 값을 학습할때 넣으면 이 카테고리들이 학습됨 names = list(request.label_map) + + # 레이블 변환기 (file_util.py/create_segmentation_train_label() 에 쓰임) + label_converter = {request.label_map[key]:idx for idx, key in enumerate(request.label_map)} + # key : 데이터에 저장된 프로젝트 카테고리 id + # value : 모델에 저장될 카테고리 id (모델에는 key의 idx 순서대로 저장될 것임) # 데이터 전처리: 학습할 디렉토리 & 데이터셋 설정 파일을 생성 process_directories(dataset_root_path, names) @@ -94,7 +98,7 @@ async def segmentation_train(request: TrainRequest): train_data, val_data = split_data(request.data, request.ratio) # 데이터 전처리: 데이터 이미지 및 레이블 다운로드 - download_data(train_data, val_data, dataset_root_path) + download_data(train_data, val_data, dataset_root_path, label_converter) # 학습 results = run_train(request, model,dataset_root_path) diff --git a/ai/app/utils/file_utils.py b/ai/app/utils/file_utils.py index 28a1523..4a9faa9 100644 --- a/ai/app/utils/file_utils.py +++ b/ai/app/utils/file_utils.py @@ -39,7 +39,7 @@ def process_directories(dataset_root_path:str, model_categories:list[str]): shutil.rmtree(os.path.join(dataset_root_path, "result")) make_yml(dataset_root_path, model_categories) -def process_image_and_label(data:TrainDataInfo, dataset_root_path:str, child_path:str): +def process_image_and_label(data:TrainDataInfo, dataset_root_path:str, child_path:str, label_converter:dict[int,int]): """이미지 저장 및 레이블 파일 생성""" # 이미지 url로부터 파일명 분리 img_name = data.image_url.split('/')[-1] @@ -60,11 +60,11 @@ def process_image_and_label(data:TrainDataInfo, dataset_root_path:str, child_pat # 레이블 -> 학습용 레이블 데이터 파싱 후 생성 if label['task_type'] == "det": - create_detection_train_label(label, label_path) + create_detection_train_label(label, label_path, label_converter) elif label["task_type"] == "seg": - create_segmentation_train_label(label, label_path) + create_segmentation_train_label(label, label_path, label_converter) -def create_detection_train_label(label:dict, label_path:str): +def create_detection_train_label(label:dict, label_path:str, label_converter:dict[int, int]): with open(label_path, "w") as train_label_txt: for shape in label["shapes"]: train_label = [] @@ -72,18 +72,18 @@ def create_detection_train_label(label:dict, label_path:str): y1 = shape["points"][0][1] x2 = shape["points"][1][0] y2 = shape["points"][1][1] - train_label.append(str(shape["group_id"])) # label Id + train_label.append(str(label_converter[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 create_segmentation_train_label(label:dict, label_path:str): +def create_segmentation_train_label(label:dict, label_path:str, label_converter:dict[int, int]): with open(label_path, "w") as train_label_txt: for shape in label["shapes"]: train_label = [] - train_label.append(str(shape["group_id"])) # label Id + train_label.append(str(label_converter[shape["group_id"]])) # label Id for x, y in shape["points"]: train_label.append(str(x / label["imageWidth"])) train_label.append(str(y / label["imageHeight"])) @@ -140,7 +140,4 @@ def process_image_and_label_in_cls(data:TrainDataInfo, dataset_root_path:str, ch label_path = os.path.join(dataset_root_path,child_path,label_name) # url로부터 이미지 다운로드 - urllib.request.urlretrieve(data.image_url, os.path.join(label_path, img_name)) - -def download_image(url, path): - urllib.request.urlretrieve(url, path) \ No newline at end of file + urllib.request.urlretrieve(data.image_url, os.path.join(label_path, img_name)) \ No newline at end of file