Merge branch 'ai/develop' of https://lab.ssafy.com/s11-s-project/S11P21S002 into ai/feat/52-ai-model-create
This commit is contained in:
김진현
commit
a5ab6c716b
@ -1,3 +1,5 @@
|
||||
import json
|
||||
|
||||
from fastapi import APIRouter, HTTPException
|
||||
from schemas.predict_request import PredictRequest
|
||||
from schemas.train_request import TrainRequest
|
||||
@ -6,12 +8,28 @@ from app.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 fastapi.responses import FileResponse
|
||||
from utils.websocket_utils import WebSocketClient
|
||||
import asyncio
|
||||
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
|
||||
@router.post("/detection", response_model=List[PredictResponse])
|
||||
def predict(request: PredictRequest):
|
||||
async def 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:
|
||||
await ws_client.connect()
|
||||
|
||||
# 모델 로드
|
||||
try:
|
||||
@ -21,13 +39,10 @@ def predict(request: PredictRequest):
|
||||
|
||||
# 추론
|
||||
results = []
|
||||
total_images = len(request.image_list)
|
||||
for idx, image in enumerate(request.image_list):
|
||||
try:
|
||||
for image in request.image_list:
|
||||
# URL에서 이미지를 메모리로 로드 TODO: 추후 메모리에 할지 어떻게 해야할지 or 병렬 처리 고민
|
||||
# response = requests.get(image.image_url)
|
||||
|
||||
# 이미지 데이터를 메모리로 로드
|
||||
# img = Image.open(io.BytesIO(response.content))
|
||||
|
||||
predict_results = model.predict(
|
||||
source=image.image_url,
|
||||
@ -35,11 +50,49 @@ def predict(request: PredictRequest):
|
||||
conf=request.conf_threshold,
|
||||
classes=request.classes
|
||||
)
|
||||
results.append(predict_results[0])
|
||||
# 예측 결과 처리
|
||||
result = predict_results[0]
|
||||
label_data = LabelData(
|
||||
version=version,
|
||||
task_type="det",
|
||||
shapes=[
|
||||
{
|
||||
"label": summary['name'],
|
||||
"color": "#ff0000",
|
||||
"points": [
|
||||
[summary['box']['x1'], summary['box']['y1']],
|
||||
[summary['box']['x2'], summary['box']['y2']]
|
||||
],
|
||||
"group_id": 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.dict()
|
||||
}
|
||||
|
||||
await ws_client.send_message("/app/ai/predict/progress", json.dumps(message))
|
||||
|
||||
# 메모리에서 이미지 객체 해제
|
||||
# img.close()
|
||||
# del img
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=500, detail="model predict exception: " + str(e))
|
||||
|
||||
@ -76,11 +129,20 @@ def predict(request: PredictRequest):
|
||||
})
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=500, detail="label parsing 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")
|
||||
|
||||
finally:
|
||||
if ws_client.is_connected():
|
||||
await ws_client.close()
|
||||
|
||||
|
||||
@router.post("/detection/train")
|
||||
def train(request: TrainRequest):
|
||||
async def train(request: TrainRequest):
|
||||
# 데이터셋 루트 경로 얻기
|
||||
dataset_root_path = get_dataset_root_path(request.project_id)
|
||||
|
||||
@ -90,16 +152,40 @@ def train(request: TrainRequest):
|
||||
# 학습 데이터 분류
|
||||
train_data, val_data = split_data(request.data, request.ratio, request.seed)
|
||||
|
||||
# Spring 서버의 WebSocket URL
|
||||
# TODO: 배포시에 변경
|
||||
spring_server_ws_url = f"ws://localhost:8080/ws"
|
||||
|
||||
# WebSocketClient 인스턴스 생성
|
||||
ws_client = WebSocketClient(spring_server_ws_url)
|
||||
|
||||
|
||||
try:
|
||||
await ws_client.connect()
|
||||
|
||||
# 학습 데이터 처리
|
||||
for data in train_data:
|
||||
process_image_and_label(data, dataset_root_path, "train")
|
||||
total_data = len(train_data)
|
||||
for idx, data in enumerate(train_data):
|
||||
# TODO: 비동기면 await 연결
|
||||
# process_image_and_label(data, dataset_root_path, "train")
|
||||
|
||||
# 진행률 계산
|
||||
progress = (idx + 1) / total_data * 100
|
||||
|
||||
await ws_client.send_message("/app/ai/train/progress", f"학습 데이터 처리 중 {request.project_id}: {progress:.2f}% 완료")
|
||||
|
||||
# 검증 데이터 처리
|
||||
for data in val_data:
|
||||
process_image_and_label(data, dataset_root_path, "val")
|
||||
total_val_data = len(val_data)
|
||||
for idx, data in enumerate(val_data):
|
||||
# TODO: 비동기면 await 연결
|
||||
# process_image_and_label(data, dataset_root_path, "val")
|
||||
|
||||
# 진행률 계산
|
||||
progress = (idx + 1) / total_val_data * 100
|
||||
# 웹소켓으로 메시지 전송 (필요할 경우 추가)
|
||||
await ws_client.send_message("/app/ai/val/progress", f"검증 데이터 처리 중 {request.project_id}: {progress:.2f}% 완료")
|
||||
|
||||
model = load_detection_model("test-data/model/best.pt")
|
||||
|
||||
model.train(
|
||||
data=join_path(dataset_root_path, "dataset.yaml"),
|
||||
name=join_path(dataset_root_path, "result"),
|
||||
@ -107,4 +193,18 @@ def train(request: TrainRequest):
|
||||
batch=request.batch,
|
||||
)
|
||||
|
||||
return FileResponse(path=join_path(dataset_root_path, "result", "weights", "best.pt"), filename="best.pt", media_type="application/octet-stream")
|
||||
# return FileResponse(path=join_path(dataset_root_path, "result", "weights", "best.pt"), filename="best.pt", media_type="application/octet-stream")
|
||||
|
||||
return {"status": "Training completed successfully"}
|
||||
|
||||
except Exception as e:
|
||||
print(f"Training process failed: {str(e)}")
|
||||
raise HTTPException(status_code=500, detail="Training process failed")
|
||||
|
||||
finally:
|
||||
if ws_client.is_connected():
|
||||
await ws_client.close()
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
@ -6,7 +6,7 @@ from ultralytics.nn.tasks import DetectionModel, SegmentationModel
|
||||
import os
|
||||
import torch
|
||||
|
||||
def load_detection_model(model_path: str = os.path.join("test-data","model","initial.pt"), device:str ="auto"):
|
||||
def load_detection_model(model_path: str = os.path.join("test-data","model","yolov8n.pt"), device:str ="auto"):
|
||||
"""
|
||||
지정된 경로에서 YOLO 모델을 로드합니다.
|
||||
|
||||
|
||||
36
ai/app/utils/websocket_utils.py
Normal file
36
ai/app/utils/websocket_utils.py
Normal file
@ -0,0 +1,36 @@
|
||||
import websockets
|
||||
|
||||
class WebSocketClient:
|
||||
def __init__(self, url: str):
|
||||
self.url = url
|
||||
self.websocket = None
|
||||
|
||||
async def connect(self):
|
||||
try:
|
||||
self.websocket = await websockets.connect(self.url)
|
||||
print(f"Connected to WebSocket at {self.url}")
|
||||
except Exception as e:
|
||||
print(f"Failed to connect to WebSocket: {str(e)}")
|
||||
|
||||
async def send_message(self, destination: str, message: str):
|
||||
try:
|
||||
if self.websocket is not None:
|
||||
# STOMP 형식의 메시지를 전송
|
||||
await self.websocket.send(f"SEND\ndestination:{destination}\n\n{message}\u0000")
|
||||
print(f"Sent message to {destination}: {message}")
|
||||
else:
|
||||
print("WebSocket is not connected. Unable to send message.")
|
||||
except Exception as e:
|
||||
print(f"Failed to send message: {str(e)}")
|
||||
return
|
||||
|
||||
async def close(self):
|
||||
try:
|
||||
if self.websocket is not None:
|
||||
await self.websocket.close()
|
||||
print("WebSocket connection closed.")
|
||||
except Exception as e:
|
||||
print(f"Failed to close WebSocket connection: {str(e)}")
|
||||
|
||||
def is_connected(self):
|
||||
return self.websocket is not None and self.websocket.open
|
||||
@ -17,3 +17,4 @@ dependencies:
|
||||
- boto3
|
||||
- python-dotenv
|
||||
- locust
|
||||
- websockets
|
||||
|
||||
Loading…
Reference in New Issue
Block a user