jhynsoo/worlabel
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Feat: segmentation 프로젝트 명세 변경에 따른 수정

Browse Source
This commit is contained in:
김진현 2024-09-27 12:03:26 +09:00
parent 0d4000b879
commit 46c971494e
3 changed files with 46 additions and 87 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
ai/app/api/yolo/detection.py
View File

@ -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

129
ai/app/api/yolo/segmentation.py
View File

@ -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 # 남은 시간(초)

1
ai/app/schemas/train_report_data.py
View File

@ -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