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

Feat: Detection prediction 웹소켓 연결 실패시에도 실행되도록 구현

Browse Source
This commit is contained in:
김진현 2024-09-19 16:49:11 +09:00
parent 47288e6bab
commit 9ce4c986d2
4 changed files with 70 additions and 44 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -1,48 +1,46 @@
import json
from fastapi import APIRouter, HTTPException
from fastapi import APIRouter, HTTPException, Response
from schemas.predict_request import PredictRequest
from schemas.train_request import TrainRequest
from schemas.predict_response import PredictResponse, LabelData
from services.load_model import load_detection_model
from utils.dataset_utils import split_data
from utils.file_utils import get_dataset_root_path, process_directories, process_image_and_label, join_path
from typing import List
from utils.websocket_utils import WebSocketClient
from utils.websocket_utils import WebSocketClient, WebSocketConnectionException
import asyncio
router = APIRouter()
@router.post("/predict", response_model=List[PredictResponse])
@router.post("/predict")
async def detection_predict(request: PredictRequest):
version = "0.1.0"
print("여기")
# Spring 서버의 WebSocket URL
# TODO: 배포 시 변경
spring_server_ws_url = f"ws://localhost:8080/ws"
print("여기")
# WebSocketClient 인스턴스 생성
ws_client = WebSocketClient(spring_server_ws_url)
# 모델 로드
try:
model = load_detection_model(request.path)
except Exception as e:
raise HTTPException(status_code=500, detail="load model exception: " + str(e))
# 웹소켓 연결
try:
await ws_client.connect()
# 모델 로드
try:
model = load_detection_model()
except Exception as e:
raise HTTPException(status_code=500, detail="load model exception: " + str(e))
if not ws_client.is_connected():
raise WebSocketConnectionException()
# 추론
results = []
total_images = len(request.image_list)
for idx, image in enumerate(request.image_list):
try:
# URL에서 이미지를 메모리로 로드 TODO: 추후 메모리에 할지 어떻게 해야할지 or 병렬 처리 고민
predict_results = model.predict(
source=image.image_url,
iou=request.iou_threshold,
@ -87,22 +85,34 @@ async def detection_predict(request: PredictRequest):
message = {
"project_id": request.project_id,
"progress": progress,
"result": response_item.dict()
"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 = []
try:
for (image, result) in zip(request.image_list, results):
label_data: LabelData = {
"version": version,
"task_type": "det",
"shapes": [
for image in request.image_list:
try:
predict_results = model.predict(
source=image.image_url,
iou=request.iou_threshold,
conf=request.conf_threshold,
classes=request.classes
)
# 예측 결과 처리
result = predict_results[0]
label_data = LabelData(
version=version,
task_type="det",
shapes=[
{
"label": summary['name'],
"color": "#ff0000",
@ -116,21 +126,24 @@ async def detection_predict(request: PredictRequest):
}
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.append({
"image_id": image.image_id,
"image_url": image.image_url,
"data": label_data
})
except Exception as e:
raise HTTPException(status_code=500, detail="label parsing exception: " + str(e))
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)
except Exception as e:
raise HTTPException(status_code=500, detail="model predict exception: " + str(e))
return response
except Exception as e:
print(f"Prediction process failed: {str(e)}")
raise HTTPException(status_code=500, detail="Prediction process failed")

16
ai/app/schemas/predict_request.py
View File

@ -1,14 +1,20 @@
from pydantic import BaseModel
from pydantic import BaseModel, Field
from typing import List, Optional
class ImageInfo(BaseModel):
image_id: int
image_url: str
image_url: str
class LabelCategory(BaseModel):
label_id: int
label_name: str
class PredictRequest(BaseModel):
project_id: int
image_list: List[ImageInfo]
version: Optional[str] = "latest"
conf_threshold: Optional[float] = 0.25
iou_threshold: Optional[float] = 0.45
version: str = "latest"
conf_threshold: float = 0.25
iou_threshold: float = 0.45
classes: Optional[List[int]] = None
path: Optional[str] = Field(None, alias="model_path")
label_categories: Optional[List[LabelCategory]] = None

6
ai/app/schemas/train_request.py
View File

@ -1,12 +1,12 @@
from pydantic import BaseModel
from pydantic import BaseModel, Field
from typing import List, Optional, Union
from schemas.predict_response import LabelData
from schemas.predict_request import LabelCategory
class TrainDataInfo(BaseModel):
image_url: str
label: LabelData
class TrainRequest(BaseModel):
project_id: int
data: List[TrainDataInfo]
@ -14,3 +14,5 @@ class TrainRequest(BaseModel):
ratio: float = 0.8 # 훈련/검증 분할 비율
epochs: int = 50 # 훈련 반복 횟수
batch: Union[float, int] = -1 # 훈련 batch 수[int] or GPU의 사용률 자동[float] default(-1): gpu의 60% 사용 유지
path: Optional[str] = Field(None, alias="model_path")
label_categories: Optional[List[LabelCategory]] = None # 새로운 레이블 카테고리 확인용

7
ai/app/utils/websocket_utils.py
View File

@ -1,4 +1,5 @@
import websockets
from websockets import WebSocketException
class WebSocketClient:
def __init__(self, url: str):
@ -33,4 +34,8 @@ class WebSocketClient:
print(f"Failed to close WebSocket connection: {str(e)}")
def is_connected(self):
return self.websocket is not None and self.websocket.open
return self.websocket is not None and self.websocket.open
class WebSocketConnectionException(WebSocketException):
def __init__(self, message="Failed to connect to WebSocket"):
super().__init__(message)