Merge branch 'ai/feat/classification' into 'ai/develop'
Fix: classification get_model() 파라미터 수정, 카테고리 없을 시 동작 수정 See merge request s11-s-project/S11P21S002!222
This commit is contained in:
commit
7ccd3b57d7
@ -1,8 +1,8 @@
|
|||||||
from fastapi import APIRouter, HTTPException
|
from fastapi import APIRouter, HTTPException
|
||||||
from api.yolo.detection import get_classes, run_predictions, get_random_color, split_data
|
from api.yolo.detection import run_predictions, get_random_color, split_data
|
||||||
from schemas.predict_request import PredictRequest
|
from schemas.predict_request import PredictRequest
|
||||||
from schemas.train_request import TrainRequest, TrainDataInfo
|
from schemas.train_request import TrainRequest, TrainDataInfo
|
||||||
from schemas.predict_response import PredictResponse, LabelData
|
from schemas.predict_response import PredictResponse, LabelData, Shape
|
||||||
from schemas.train_report_data import ReportData
|
from schemas.train_report_data import ReportData
|
||||||
from schemas.train_response import TrainResponse
|
from schemas.train_response import TrainResponse
|
||||||
from services.load_model import load_classification_model
|
from services.load_model import load_classification_model
|
||||||
@ -24,11 +24,8 @@ async def classification_predict(request: PredictRequest):
|
|||||||
# 이미지 데이터 정리
|
# 이미지 데이터 정리
|
||||||
url_list = list(map(lambda x:x.image_url, request.image_list))
|
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)
|
results = run_predictions(model, url_list, request, classes=[]) # classification은 classes를 무시함
|
||||||
|
|
||||||
# 추론 결과 변환
|
# 추론 결과 변환
|
||||||
response = [process_prediction_result(result, image, request.label_map) for result, image in zip(results,request.image_list)]
|
response = [process_prediction_result(result, image, request.label_map) for result, image in zip(results,request.image_list)]
|
||||||
@ -36,43 +33,53 @@ async def classification_predict(request: PredictRequest):
|
|||||||
return response
|
return response
|
||||||
|
|
||||||
# 모델 로드
|
# 모델 로드
|
||||||
def get_model(request: PredictRequest):
|
def get_model(project_id:int, model_key:str):
|
||||||
try:
|
try:
|
||||||
return load_classification_model(request.project_id, request.m_key)
|
return load_classification_model(project_id, model_key)
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
raise HTTPException(status_code=500, detail="exception in get_model(): " + str(e))
|
raise HTTPException(status_code=500, detail="exception in get_model(): " + str(e))
|
||||||
|
|
||||||
# 추론 결과 처리 함수
|
# 추론 결과 처리 함수
|
||||||
def process_prediction_result(result, image, label_map):
|
def process_prediction_result(result, image, label_map):
|
||||||
try:
|
try:
|
||||||
|
label_name = None
|
||||||
|
# top 5에 해당하는 class id 순회
|
||||||
|
for class_id in result.probs.top5:
|
||||||
|
name = result.names[class_id] # class id에 해당하는 label_name
|
||||||
|
if name in label_map: # name이 사용자 레이블 카테고리에 있을 경우
|
||||||
|
label_name = name # label_name 설정
|
||||||
|
break
|
||||||
|
|
||||||
label_data = LabelData(
|
label_data = LabelData(
|
||||||
version="0.0.0",
|
version="0.0.0",
|
||||||
task_type="cls",
|
task_type="cls",
|
||||||
shapes=[
|
shapes=[],
|
||||||
{
|
|
||||||
"label": summary['name'],
|
|
||||||
"color": get_random_color(),
|
|
||||||
"points": [
|
|
||||||
[0, 0]
|
|
||||||
],
|
|
||||||
"group_id": label_map[summary['name']],
|
|
||||||
"shape_type": "point",
|
|
||||||
"flags": {}
|
|
||||||
}
|
|
||||||
for summary in result.summary()
|
|
||||||
],
|
|
||||||
split="none",
|
split="none",
|
||||||
imageHeight=result.orig_img.shape[0],
|
imageHeight=result.orig_img.shape[0],
|
||||||
imageWidth=result.orig_img.shape[1],
|
imageWidth=result.orig_img.shape[1],
|
||||||
imageDepth=result.orig_img.shape[2]
|
imageDepth=result.orig_img.shape[2]
|
||||||
)
|
)
|
||||||
|
|
||||||
|
if label_name: # label_name을 설정한게 있다면 추가
|
||||||
|
shape = Shape(
|
||||||
|
label= label_name,
|
||||||
|
color= get_random_color(),
|
||||||
|
points= [[0.0, 0.0]],
|
||||||
|
group_id= label_map[label_name],
|
||||||
|
shape_type= 'point',
|
||||||
|
flags= {}
|
||||||
|
)
|
||||||
|
LabelData.shapes.append(shape)
|
||||||
|
|
||||||
|
return PredictResponse(
|
||||||
|
image_id=image.image_id,
|
||||||
|
data=label_data.model_dump_json()
|
||||||
|
)
|
||||||
|
except KeyError as e:
|
||||||
|
raise HTTPException(status_code=500, detail="KeyError: " + str(e))
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
raise HTTPException(status_code=500, detail="exception in process_prediction_result(): " + str(e))
|
raise HTTPException(status_code=500, detail="exception in process_prediction_result(): " + str(e))
|
||||||
|
|
||||||
return PredictResponse(
|
|
||||||
image_id=image.image_id,
|
|
||||||
data=label_data.model_dump_json()
|
|
||||||
)
|
|
||||||
|
|
||||||
@router.post("/train")
|
@router.post("/train")
|
||||||
async def classification_train(request: TrainRequest):
|
async def classification_train(request: TrainRequest):
|
||||||
|
Loading…
Reference in New Issue
Block a user