Feat: Fast API 소켓 열기

ai/app/api/yolo/detection.py

@@ -7,6 +7,8 @@ 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 handler.websocket_handler import WebSocketClient
+import asyncio
 
 router = APIRouter()
 @router.post("/detection", response_model=List[PredictResponse])
@@ -80,31 +82,60 @@ def predict(request: PredictRequest):
 
 
 @router.post("/detection/train")
-def train(request: TrainRequest):
-    # 데이터셋 루트 경로 얻기
-    dataset_root_path = get_dataset_root_path(request.project_id)
+async def train(request: TrainRequest):
+    # Spring 서버의 WebSocket URL
+    # spring_server_ws_url = f"ws://localhost:8080/ws/ai/train/progress/{request.project_id}"
+    spring_server_ws_url = f"ws://localhost:8080/ws"
     
-    # 디렉토리 생성 및 초기화
-    process_directories(dataset_root_path)
-    
-    # 학습 데이터 분류
-    train_data, val_data = split_data(request.data, request.ratio, request.seed)
-    
-    # 학습 데이터 처리
-    for data in train_data:
-        process_image_and_label(data, dataset_root_path, "train")
+    # WebSocketClient 인스턴스 생성
+    print("연결 요청 - " + spring_server_ws_url)
+    ws_client = WebSocketClient(spring_server_ws_url)
 
-    # 검증 데이터 처리
-    for data in val_data:
-        process_image_and_label(data, dataset_root_path, "val")
+    try:
+        await ws_client.connect()
 
-    model = load_detection_model("test-data/model/best.pt")
+        await ws_client.send_message("/app/ai/train/progress", f"Training started for project {request.project_id}")
 
-    model.train(
-        data=join_path(dataset_root_path,"dataset.yaml"),
-        name=join_path(dataset_root_path,"result"), 
-        epochs= request.epochs, 
-        batch=request.batch, 
-        )
+        for i in range(1, 31):
+            await ws_client.send_message("/app/ai/train/progress", f"Training progress: {i}/30")
+            await asyncio.sleep(1)
+
+        await ws_client.send_message("/app/ai/train/progress", "Training complete")
+
+        return {"status": "Training completed successfully"}
+
+    except Exception as e:
+        logging.error(f"Training process failed: {str(e)}")
+        raise HTTPException(status_code=500, detail="Training process failed")
+
+    finally:
+        await ws_client.close()
     
-    return FileResponse(path=join_path(dataset_root_path, "result", "weights", "best.pt"), filename="best.pt", media_type="application/octet-stream")
+
+    # # 데이터셋 루트 경로 얻기
+    # dataset_root_path = get_dataset_root_path(request.project_id)
+    
+    # # 디렉토리 생성 및 초기화
+    # process_directories(dataset_root_path)
+    
+    # # 학습 데이터 분류
+    # train_data, val_data = split_data(request.data, request.ratio, request.seed)
+    
+    # # 학습 데이터 처리
+    # for data in train_data:
+    #     process_image_and_label(data, dataset_root_path, "train")
+
+    # # 검증 데이터 처리
+    # for data in val_data:
+    #     process_image_and_label(data, dataset_root_path, "val")
+
+    # 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"), 
+    #     epochs= request.epochs, 
+    #     batch=request.batch, 
+    #     )
+    
+    # return FileResponse(path=join_path(dataset_root_path, "result", "weights", "best.pt"), filename="best.pt", media_type="application/octet-stream")

ai/app/handler/websocket_handler.py

@@ -0,0 +1,35 @@
+import asyncio
+import websockets
+import logging
+
+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)}")

ai/environment.yml

@@ -16,4 +16,5 @@ dependencies:
   - dill
   - boto3
   - python-dotenv
-  - locust
+  - locust
+  - websockets