Fix: 리팩토링에서 생긴 버그 수정
This commit is contained in:
김진현
parent
81b5145b41
commit
27cfd01fd7
@ -18,7 +18,7 @@ router = APIRouter()
|
||||
@router.post("/predict")
|
||||
async def detection_predict(request: PredictRequest):
|
||||
|
||||
send_slack_message(f"predict 요청{request}", status="success")
|
||||
# send_slack_message(f"predict 요청{request}", status="success")
|
||||
|
||||
# 모델 로드
|
||||
model = get_model(request)
|
||||
@ -33,6 +33,8 @@ async def detection_predict(request: PredictRequest):
|
||||
# 추론
|
||||
try:
|
||||
results = run_predictions(model, url_list, request, classes)
|
||||
print(len(results))
|
||||
print(len(request.image_list))
|
||||
response = [process_prediction_result(result, image, request.label_map) for result, image in zip(results,request.image_list)]
|
||||
return response
|
||||
|
||||
@ -51,13 +53,12 @@ def get_model(request: PredictRequest):
|
||||
# 추론 실행 함수
|
||||
def run_predictions(model, image, request, classes):
|
||||
try:
|
||||
predict_results = model.predict(
|
||||
source=image,
|
||||
iou=request.iou_threshold,
|
||||
conf=request.conf_threshold,
|
||||
classes=classes
|
||||
return model.predict(
|
||||
source=image,
|
||||
iou=request.iou_threshold,
|
||||
conf=request.conf_threshold,
|
||||
classes=classes
|
||||
)
|
||||
return predict_results[0]
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=500, detail="model predict exception: " + str(e))
|
||||
|
||||
|
||||
@ -11,7 +11,7 @@ def predict(request: PredictRequest):
|
||||
|
||||
# 모델 로드
|
||||
try:
|
||||
model = load_segmentation_model()
|
||||
model = load_segmentation_model(request.project_id, request.m_key)
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=500, detail="load model exception: "+str(e))
|
||||
|
||||
|
||||
@ -10,7 +10,7 @@ app.include_router(yolo_detection_router, prefix="/api/detection", tags=["Detect
|
||||
app.include_router(yolo_segmentation_router, prefix="/api/segmentation", tags=["Segmentation"])
|
||||
app.include_router(yolo_model_router, prefix="/api/model", tags=["Model"])
|
||||
|
||||
# 애플리케이션 실행
|
||||
if __name__ == "__main__":
|
||||
import uvicorn
|
||||
uvicorn.run("main:app", reload=True)
|
||||
# # 애플리케이션 실행
|
||||
# if __name__ == "__main__":
|
||||
# import uvicorn
|
||||
# uvicorn.run("main:app", reload=True)
|
||||
|
||||
@ -8,9 +8,9 @@ class TrainDataInfo(BaseModel):
|
||||
|
||||
class TrainRequest(BaseModel):
|
||||
project_id: int
|
||||
m_key: Optional[str] = Field(None, alias="model_key")
|
||||
m_id: Optional[str] = Field(None, alias="model_id") # 학습 중 에포크 결과를 보낼때 model_id를 보냄
|
||||
label_map: dict[int, int] = Field(None, description="모델 레이블 카테고리 idx: 프로젝트 레이블 카테고리 idx , None 일경우 레이블 데이터(프로젝트 레이블)의 idx로 학습")
|
||||
m_key: str = Field("yolo8", alias="model_key")
|
||||
m_id: int = Field(..., alias="model_id") # 학습 중 에포크 결과를 보낼때 model_id를 보냄
|
||||
label_map: dict[int, int] = Field({}, description="모델 레이블 카테고리 idx: 프로젝트 레이블 카테고리 idx , None 일경우 레이블 데이터(프로젝트 레이블)의 idx로 학습")
|
||||
data: List[TrainDataInfo]
|
||||
ratio: float = 0.8 # 훈련/검증 분할 비율
|
||||
|
||||
|
||||
@ -1,8 +1,6 @@
|
||||
# ai_service.py
|
||||
|
||||
from ultralytics import YOLO # Ultralytics YOLO 모델을 가져오기
|
||||
from ultralytics.models.yolo.model import YOLO as YOLO_Model
|
||||
from ultralytics.nn.tasks import DetectionModel, SegmentationModel
|
||||
import os
|
||||
import torch
|
||||
|
||||
@ -12,7 +10,7 @@ def load_detection_model(project_id:int, model_key:str):
|
||||
if model_key in default_model_map:
|
||||
model = YOLO(default_model_map[model_key])
|
||||
else:
|
||||
model = load_model(model_path=os.path.join("projects",str(project_id),"models",model_key))
|
||||
model = load_model(model_path=os.path.join("projects",str(project_id),"models", model_key))
|
||||
|
||||
# Detection 모델인지 검증
|
||||
if model.task != "detect":
|
||||
|
||||
@ -40,7 +40,6 @@ def process_directories(dataset_root_path:str, names:list[str]):
|
||||
make_yml(dataset_root_path, names)
|
||||
|
||||
def process_image_and_label(data:TrainDataInfo, dataset_root_path:str, child_path:str, label_map:dict[int, int]|None):
|
||||
|
||||
"""이미지 저장 및 레이블 파일 생성"""
|
||||
# 이미지 url로부터 파일명 분리
|
||||
img_name = data.image_url.split('/')[-1]
|
||||
@ -64,17 +63,17 @@ def process_image_and_label(data:TrainDataInfo, dataset_root_path:str, child_pat
|
||||
|
||||
def create_detection_train_label(label:LabelData, label_path:str, label_map:dict[int, int]|None):
|
||||
with open(label_path, "w") as train_label_txt:
|
||||
for shape in label.shapes:
|
||||
for shape in label["shapes"]:
|
||||
train_label = []
|
||||
x1 = shape.points[0][0]
|
||||
y1 = shape.points[0][1]
|
||||
x2 = shape.points[1][0]
|
||||
y2 = shape.points[1][1]
|
||||
train_label.append(str(label_map[shape.group_id]) if label_map else str(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 )) # 높이
|
||||
x1 = shape["points"][0][0]
|
||||
y1 = shape["points"][0][1]
|
||||
x2 = shape["points"][1][0]
|
||||
y2 = shape["points"][1][1]
|
||||
train_label.append(str(label_map[shape["group_id"]]) if label_map else str(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 join_path(path, *paths):
|
||||
|
||||
Loading…
Reference in New Issue
Block a user