Feat: 백엔드 API 연동
This commit is contained in:
parent
c6ef8ea5fc
commit
a827a967c7
2
ai/.gitignore
vendored
2
ai/.gitignore
vendored
@ -38,3 +38,5 @@ test-data/
|
|||||||
resources/
|
resources/
|
||||||
datasets/
|
datasets/
|
||||||
*.pt
|
*.pt
|
||||||
|
|
||||||
|
*.jpg
|
@ -23,139 +23,89 @@ async def detection_predict(request: PredictRequest):
|
|||||||
ws_client = WebSocketClient(spring_server_ws_url)
|
ws_client = WebSocketClient(spring_server_ws_url)
|
||||||
|
|
||||||
# 모델 로드
|
# 모델 로드
|
||||||
try:
|
model = load_model(request)
|
||||||
model_path = request.m_key and get_model_path(request.project_id, request.m_key)
|
|
||||||
model = load_detection_model(model_path=model_path)
|
|
||||||
except Exception as e:
|
|
||||||
raise HTTPException(status_code=500, detail="load model exception: " + str(e))
|
|
||||||
|
|
||||||
# 모델 레이블 카테고리 연결
|
# 모델 레이블 카테고리 연결
|
||||||
classes = None
|
classes = list(request.label_map) if request.label_map else None
|
||||||
if request.label_map:
|
|
||||||
classes = list(request.label_map)
|
|
||||||
|
|
||||||
|
# 결과를 저장할 리스트
|
||||||
|
response = []
|
||||||
|
|
||||||
# 웹소켓 연결
|
# 웹소켓 연결
|
||||||
try:
|
await connect_to_websocket(ws_client)
|
||||||
await ws_client.connect()
|
|
||||||
if not ws_client.is_connected():
|
|
||||||
raise WebSocketConnectionException()
|
|
||||||
|
|
||||||
# 추론
|
# 추론
|
||||||
total_images = len(request.image_list)
|
try:
|
||||||
for idx, image in enumerate(request.image_list):
|
for idx, image in enumerate(request.image_list):
|
||||||
try:
|
result = run_predictions(model, image, request, classes)
|
||||||
# URL에서 이미지를 메모리로 로드 TODO: 추후 메모리에 할지 어떻게 해야할지 or 병렬 처리 고민
|
response_item = process_prediction_result(result, request, image)
|
||||||
predict_results = model.predict(
|
|
||||||
source=image.image_url,
|
|
||||||
iou=request.iou_threshold,
|
|
||||||
conf=request.conf_threshold,
|
|
||||||
classes=classes
|
|
||||||
)
|
|
||||||
# 예측 결과 처리
|
|
||||||
result = predict_results[0]
|
|
||||||
label_data = LabelData(
|
|
||||||
version="0.0.0",
|
|
||||||
task_type="det",
|
|
||||||
shapes=[
|
|
||||||
{
|
|
||||||
"label": summary['name'],
|
|
||||||
"color": "#ff0000",
|
|
||||||
"points": [
|
|
||||||
[summary['box']['x1'], summary['box']['y1']],
|
|
||||||
[summary['box']['x2'], summary['box']['y2']]
|
|
||||||
],
|
|
||||||
"group_id": request.label_map[summary['class']] if request.label_map else summary['class'],
|
|
||||||
"shape_type": "rectangle",
|
|
||||||
"flags": {}
|
|
||||||
}
|
|
||||||
for summary in result.summary()
|
|
||||||
],
|
|
||||||
split="none",
|
|
||||||
imageHeight=result.orig_img.shape[0],
|
|
||||||
imageWidth=result.orig_img.shape[1],
|
|
||||||
imageDepth=result.orig_img.shape[2]
|
|
||||||
)
|
|
||||||
|
|
||||||
response_item = PredictResponse(
|
|
||||||
image_id=image.image_id,
|
|
||||||
image_url=image.image_url,
|
|
||||||
data=label_data
|
|
||||||
)
|
|
||||||
|
|
||||||
# 진행률 계산
|
|
||||||
progress = (idx + 1) / total_images * 100
|
|
||||||
|
|
||||||
# 웹소켓으로 예측 결과와 진행률 전송
|
|
||||||
message = {
|
|
||||||
"project_id": request.project_id,
|
|
||||||
"progress": progress,
|
|
||||||
"result": response_item.model_dump()
|
|
||||||
}
|
|
||||||
|
|
||||||
await ws_client.send_message("/app/ai/predict/progress", json.dumps(message))
|
|
||||||
|
|
||||||
except Exception as e:
|
|
||||||
raise HTTPException(status_code=500, detail="model predict exception: " + str(e))
|
|
||||||
return Response(status_code=204)
|
|
||||||
|
|
||||||
# 웹소켓 연결 안된 경우
|
|
||||||
except WebSocketConnectionException as e:
|
|
||||||
# 추론
|
|
||||||
response = []
|
|
||||||
for image in request.image_list:
|
|
||||||
try:
|
|
||||||
predict_results = model.predict(
|
|
||||||
source=image.image_url,
|
|
||||||
iou=request.iou_threshold,
|
|
||||||
conf=request.conf_threshold,
|
|
||||||
classes=classes
|
|
||||||
)
|
|
||||||
# 예측 결과 처리
|
|
||||||
result = predict_results[0]
|
|
||||||
label_data = LabelData(
|
|
||||||
version="0.0.0",
|
|
||||||
task_type="det",
|
|
||||||
shapes=[
|
|
||||||
{
|
|
||||||
"label": summary['name'],
|
|
||||||
"color": "#ff0000",
|
|
||||||
"points": [
|
|
||||||
[summary['box']['x1'], summary['box']['y1']],
|
|
||||||
[summary['box']['x2'], summary['box']['y2']]
|
|
||||||
],
|
|
||||||
"group_id": request.label_map[summary['class']] if request.label_map else summary['class'],
|
|
||||||
"shape_type": "rectangle",
|
|
||||||
"flags": {}
|
|
||||||
}
|
|
||||||
for summary in result.summary()
|
|
||||||
],
|
|
||||||
split="none",
|
|
||||||
imageHeight=result.orig_img.shape[0],
|
|
||||||
imageWidth=result.orig_img.shape[1],
|
|
||||||
imageDepth=result.orig_img.shape[2]
|
|
||||||
)
|
|
||||||
|
|
||||||
response_item = PredictResponse(
|
|
||||||
image_id=image.image_id,
|
|
||||||
image_url=image.image_url,
|
|
||||||
data=label_data
|
|
||||||
)
|
|
||||||
|
|
||||||
response.append(response_item)
|
response.append(response_item)
|
||||||
|
|
||||||
except Exception as e:
|
|
||||||
raise HTTPException(status_code=500, detail="model predict exception: " + str(e))
|
|
||||||
|
|
||||||
return response
|
return response
|
||||||
except Exception as e:
|
|
||||||
print(f"Prediction process failed: {str(e)}")
|
|
||||||
raise HTTPException(status_code=500, detail="Prediction process failed")
|
|
||||||
|
|
||||||
finally:
|
finally:
|
||||||
if ws_client.is_connected():
|
if ws_client.is_connected():
|
||||||
await ws_client.close()
|
await ws_client.close()
|
||||||
|
|
||||||
|
# 모델 로드
|
||||||
|
def load_model(request: PredictRequest):
|
||||||
|
try:
|
||||||
|
model_path = request.m_key and get_model_path(request.project_id, request.m_key)
|
||||||
|
return load_detection_model(model_path=model_path)
|
||||||
|
except Exception as e:
|
||||||
|
raise HTTPException(status_code=500, detail="load model exception: " + str(e))
|
||||||
|
|
||||||
|
# 웹소켓 연결
|
||||||
|
async def connect_to_websocket(ws_client):
|
||||||
|
try:
|
||||||
|
await ws_client.connect()
|
||||||
|
if not ws_client.is_connected():
|
||||||
|
raise WebSocketConnectionException("웹 소켓 연결 실패")
|
||||||
|
except Exception as e:
|
||||||
|
raise HTTPException(status_code=500, detail="websocket connect failed: " + str(e))
|
||||||
|
|
||||||
|
# 추론 실행 함수
|
||||||
|
def run_predictions(model, image, request, classes):
|
||||||
|
try:
|
||||||
|
predict_results = model.predict(
|
||||||
|
source=image.image_url,
|
||||||
|
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))
|
||||||
|
|
||||||
|
# 추론 결과 처리 함수
|
||||||
|
def process_prediction_result(result, request, image):
|
||||||
|
label_data = LabelData(
|
||||||
|
version="0.0.0",
|
||||||
|
task_type="det",
|
||||||
|
shapes=[
|
||||||
|
{
|
||||||
|
"label": summary['name'],
|
||||||
|
"color": "#ff0000",
|
||||||
|
"points": [
|
||||||
|
[summary['box']['x1'], summary['box']['y1']],
|
||||||
|
[summary['box']['x2'], summary['box']['y2']]
|
||||||
|
],
|
||||||
|
"group_id": request.label_map[summary['class']] if request.label_map else summary['class'],
|
||||||
|
"shape_type": "rectangle",
|
||||||
|
"flags": {}
|
||||||
|
}
|
||||||
|
for summary in result.summary()
|
||||||
|
],
|
||||||
|
split="none",
|
||||||
|
imageHeight=result.orig_img.shape[0],
|
||||||
|
imageWidth=result.orig_img.shape[1],
|
||||||
|
imageDepth=result.orig_img.shape[2]
|
||||||
|
)
|
||||||
|
|
||||||
|
return PredictResponse(
|
||||||
|
image_id=image.image_id,
|
||||||
|
data=json.dumps(label_data.dict())
|
||||||
|
)
|
||||||
|
|
||||||
@router.post("/train")
|
@router.post("/train")
|
||||||
async def detection_train(request: TrainRequest):
|
async def detection_train(request: TrainRequest):
|
||||||
|
@ -20,5 +20,4 @@ class LabelData(BaseModel):
|
|||||||
|
|
||||||
class PredictResponse(BaseModel):
|
class PredictResponse(BaseModel):
|
||||||
image_id: int
|
image_id: int
|
||||||
image_url: str
|
data: str
|
||||||
data: LabelData
|
|
Loading…
Reference in New Issue
Block a user