Merge branch 'ai/develop' into 'master'

Ai/develop

See merge request s11-s-project/S11P21S002!312

ai/.gitignore

@@ -0,0 +1,42 @@
+# Python 기본
+__pycache__/
+*.py[cod]
+*.pyo
+*.pyd
+
+# 환경 설정 파일
+.env
+*.env
+
+# 패키지 디렉토리
+.venv/
+venv/
+env/
+
+# 빌드 디렉토리
+build/
+dist/
+*.egg-info/
+
+# 로그 파일
+*.log
+
+# Jupyter Notebook 체크포인트
+.ipynb_checkpoints
+
+# IDE 관련 파일
+.vscode/
+.idea/
+
+# MacOS 관련 파일
+.DS_Store
+
+# 테스트 파일
+test-data/
+
+# 리소스
+resources/
+datasets/
+*.pt
+
+*.jpg

ai/README.md

@@ -0,0 +1,30 @@
+# FastAPI를 이용한 AI 모델 관련 API
+
+## conda 환경 세팅
+```bash
+conda env create -f environment.yml
+conda activate worlabel_ai_env
+```
+
+## FastAPI Project 구조
+
+### app/api
+- api 호출 라우터 정의
+
+### app/schemas
+- api의 request/response 등 Pydantic 모델 정의
+
+### app/services
+- AI 관련 패키지를 이용하는 메서드 정의
+
+### app/utils
+- 프로젝트 전역에서 이용하는 formatter 등 정의
+
+### resources/models
+- yolo 기본 모델 6종(default/pretrained, det/seg/cls) 저장
+
+### resources/projects/{project_id}/models
+- 프로젝트별 ai 모델 저장
+
+### resources/datasets
+- 훈련 데이터셋 저장

ai/app/api/yolo/classfication.py

@@ -0,0 +1,199 @@
+from fastapi import APIRouter, HTTPException
+from api.yolo.detection import run_predictions, get_random_color, split_data
+from schemas.predict_request import PredictRequest
+from schemas.train_request import TrainRequest, TrainDataInfo
+from schemas.predict_response import PredictResponse, LabelData, Shape
+from schemas.train_report_data import ReportData
+from schemas.train_response import TrainResponse
+from services.load_model import load_classification_model
+from services.create_model import save_model
+from utils.file_utils import get_dataset_root_path, process_directories_in_cls, process_image_and_label_in_cls, join_path
+from utils.slackMessage import send_slack_message
+from utils.api_utils import send_data_call_api
+
+router = APIRouter()
+
+@router.post("/predict")
+async def classification_predict(request: PredictRequest):
+
+    send_slack_message(f"predict 요청: {request}", status="success")
+
+    # 모델 로드
+    model = get_model(request.project_id, request.m_key)
+
+    # 이미지 데이터 정리
+    url_list = list(map(lambda x:x.image_url, request.image_list))
+
+    # 추론
+    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)]
+    send_slack_message(f"predict 성공{response}", status="success")
+    return response
+
+# 모델 로드
+def get_model(project_id:int, model_key:str):
+    try:
+        return load_classification_model(project_id, model_key)
+    except Exception as e:
+        raise HTTPException(status_code=500, detail="exception in get_model(): " + str(e)) 
+
+# 추론 결과 처리 함수
+def process_prediction_result(result, image, label_map):
+    try:
+        shapes = []
+        # top 5에 해당하는 class id 순회
+        for class_id in result.probs.top5:
+            label_name = result.names[class_id]   # class id에 해당하는 label_name
+            if label_name in label_map:           # name이 사용자 레이블 카테고리에 있을 경우
+                shapes = [
+                    Shape(
+                        label=label_name,
+                        color=get_random_color(),
+                        points=[[0.0, 0.0]],
+                        group_id=label_map[label_name],
+                        shape_type='point',
+                        flags={}
+                    )
+                ]          # label_name 설정
+                break
+
+        label_data = LabelData(
+            version="0.0.0",
+            task_type="cls",
+            shapes=shapes,
+            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=label_data.model_dump_json()
+        )
+    
+    except KeyError as e:
+        raise HTTPException(status_code=500, detail="KeyError: " + str(e))
+    except Exception as e:
+        raise HTTPException(status_code=500, detail="exception in process_prediction_result(): " + str(e))
+
+
+@router.post("/train")
+async def classification_train(request: TrainRequest):
+
+    send_slack_message(f"train 요청{request}", status="success")
+
+    # 데이터셋 루트 경로 얻기 (프로젝트 id 기반)
+    dataset_root_path = get_dataset_root_path(request.project_id)
+
+    # 모델 로드
+    model = get_model(request.project_id, request.m_key)
+
+    # 이 값을 학습할때 넣으면 이 카테고리들이 학습됨
+    names = list(request.label_map)
+    
+    # 데이터 전처리: 학습할 디렉토리 & 데이터셋 설정 파일을 생성
+    process_directories_in_cls(dataset_root_path, names)
+
+    # 데이터 전처리: 데이터를 학습데이터와 테스트 데이터로 분류
+    train_data, test_data = split_data(request.data, request.ratio)
+
+    # 데이터 전처리: 데이터 이미지 및 레이블 다운로드
+    download_data(train_data, test_data, dataset_root_path)
+
+    # 학습
+    results = run_train(request, model,dataset_root_path)
+
+    # best 모델 저장
+    model_key = save_model(project_id=request.project_id, path=join_path(dataset_root_path, "result", "weights", "best.pt"))
+    
+    result = results.results_dict
+
+    response = TrainResponse(
+        modelKey=model_key,
+        precision= 0,
+        recall= 0,
+        mAP50= 0,
+        mAP5095= 0,
+        accuracy=result["metrics/accuracy_top1"],
+        fitness= result["fitness"]
+    )
+
+    send_slack_message(f"train 성공{response}", status="success")
+            
+    return response
+
+def download_data(train_data:list[TrainDataInfo], test_data:list[TrainDataInfo], dataset_root_path:str):
+    try:
+        for data in train_data:
+            process_image_and_label_in_cls(data, dataset_root_path, "train")
+
+        for data in test_data:
+            process_image_and_label_in_cls(data, dataset_root_path, "test")
+    except Exception as e:
+        raise HTTPException(status_code=500, detail="exception in download_data(): " + str(e))
+
+    
+def run_train(request, model, dataset_root_path):
+    try:
+        # 데이터 전송 콜백함수
+        def send_data(trainer):
+            try:
+                # 첫번째 epoch는 스킵
+                if trainer.epoch == 0:
+                    return
+
+                # 남은 시간 계산(초)
+                left_epochs = trainer.epochs - trainer.epoch
+                left_seconds = left_epochs * trainer.epoch_time
+
+                # 로스 box_loss, cls_loss, dfl_loss
+                loss = trainer.label_loss_items(loss_items=trainer.loss_items)
+                data = ReportData(
+                    epoch=trainer.epoch,            # 현재 에포크
+                    total_epochs=trainer.epochs,    # 전체 에포크
+                    seg_loss=0,                     # seg loss
+                    box_loss=0,                     # box loss
+                    cls_loss=loss["train/loss"],    # cls loss
+                    dfl_loss=0,                     # dfl loss
+                    fitness=trainer.fitness,        # 적합도
+                    epoch_time=trainer.epoch_time,  # 지난 에포크 걸린 시간 (에포크 시작 기준으로 결정)
+                    left_seconds=left_seconds       # 남은 시간(초)
+                )
+                # 데이터 전송
+                send_data_call_api(request.project_id, request.m_id, data)
+            except Exception as e:
+                print(f"Exception in send_data(): {e}")
+
+        # 콜백 등록
+        model.add_callback("on_train_epoch_start", send_data)
+
+        # 학습 실행
+        try:
+            results = model.train(
+                data=dataset_root_path,
+                name=join_path(dataset_root_path, "result"),
+                epochs=request.epochs,
+                batch=request.batch,
+                lr0=request.lr0,
+                lrf=request.lrf,
+                optimizer=request.optimizer,
+                patience=0
+            )
+        finally:
+            # 콜백 해제 및 자원 해제
+            model.reset_callbacks()
+        # 마지막 에포크 전송 
+        model.trainer.epoch += 1
+        send_data(model.trainer)
+
+        return results
+
+    except HTTPException as e:
+        raise e # HTTP 예외를 다시 발생
+    except Exception as e:
+        raise HTTPException(status_code=500, detail="exception in run_train(): "+str(e))
+
+

ai/app/api/yolo/detection.py

@@ -0,0 +1,304 @@
+import os
+import time
+
+import psutil
+from fastapi import APIRouter, HTTPException
+from schemas.predict_request import PredictRequest
+from schemas.train_request import TrainRequest, TrainDataInfo
+from schemas.predict_response import PredictResponse, LabelData, Shape
+from schemas.train_report_data import ReportData
+from schemas.train_response import TrainResponse
+from services.load_model import load_detection_model
+from services.create_model import save_model
+from utils.file_utils import get_dataset_root_path, process_directories, join_path, process_image_and_label
+from utils.slackMessage import send_slack_message
+from utils.api_utils import send_data_call_api
+import random, torch
+
+
+router = APIRouter()
+
+@router.post("/predict")
+async def detection_predict(request: PredictRequest):
+    project_id = request.project_id
+    send_slack_message(f"Detection predict 요청 (projectId: {project_id})", status="success")
+
+    # 모델 로드
+    start_time = time.time()
+    send_slack_message(f"모델 로드 중 (projectId: {project_id})...", status="success")
+    model = get_model(request.project_id, request.m_key)
+    send_slack_message(f"모델 로드 완료 (projectId: {project_id}). 걸린 시간: {time.time() - start_time:.2f} 초", status="success")
+
+    # 이미지 데이터 정리
+    start_time = time.time()
+    url_list = list(map(lambda x: x.image_url, request.image_list))
+    send_slack_message(f"이미지 데이터 정리 완료 (projectId: {project_id}). 걸린 시간: {time.time() - start_time:.2f} 초",
+                       status="success")
+
+    # 이 값을 모델에 입력하면 해당하는 클래스 id만 출력됨
+    classes = get_classes(request.label_map, model.names)
+
+    # 추론
+    start_time = time.time()
+    send_slack_message(f"추론 시작 (projectId: {project_id})...", status="success")
+    results = run_predictions(model, url_list, request, classes)
+    send_slack_message(f"추론 완료 (projectId: {project_id}). 걸린 시간: {time.time() - start_time:.2f} 초", status="success")
+
+    # 추론 결과 변환
+    start_time = time.time()
+    response = [process_prediction_result(result, image, request.label_map) for result, image in
+                zip(results, request.image_list)]
+    send_slack_message(f"추론 결과 변환 완료 (projectId: {project_id}). 걸린 시간: {time.time() - start_time:.2f} 초",
+                       status="success")
+
+    send_slack_message(f"Detection predict 성공 (projectId: {project_id}) {len(response)}", status="success")
+
+    return response
+
+# 모델 로드
+def get_model(project_id, model_key):
+    try:
+        return load_detection_model(project_id, model_key)
+    except Exception as e:
+        raise HTTPException(status_code=500, detail="exception in get_model(): " + str(e))
+
+# 모델의 레이블로부터 label_map의 key에 존재하는 값의 id만 가져오기
+def get_classes(label_map:dict[str: int], model_names: dict[int, str]):
+    try:
+        return [id for id, name in model_names.items() if name in label_map]
+    except Exception as e:
+        raise HTTPException(status_code=500, detail="exception in get_classes(): " + str(e))
+
+# 추론 실행 함수
+def run_predictions(model, image, request, classes):
+    try:
+        with torch.no_grad():
+            results = []
+            for img in image:
+                result = model.predict(
+                    source=[img],
+                    iou=request.iou_threshold,
+                    conf=request.conf_threshold,
+                    classes=classes
+                )
+                results += result
+        return results
+    except Exception as e:
+        raise HTTPException(status_code=500, detail="exception in run_predictions: " + str(e))
+
+# 추론 결과 처리 함수
+def process_prediction_result(result, image, label_map):
+    try:
+        label_data = LabelData(
+            version="0.0.0",
+            task_type="det",
+            shapes=[
+                Shape(
+                    label= summary['name'],
+                    color= get_random_color(),
+                    points= [
+                        [summary['box']['x1'], summary['box']['y1']],
+                        [summary['box']['x2'], summary['box']['y2']]
+                    ],
+                    group_id= label_map[summary['name']],
+                    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]
+        )
+    except KeyError as e:
+        raise HTTPException(status_code=500, detail="KeyError: " + str(e))
+    except Exception as 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()
+    )
+
+def get_random_color():
+    random_number = random.randint(0, 0xFFFFFF)
+    return f"#{random_number:06X}"
+
+@router.post("/train", response_model=TrainResponse)
+async def detection_train(request: TrainRequest):
+
+    send_slack_message(f"Detection train 요청 projectId : {request.project_id}, 이미지 개수:{len(request.data)}", status="success")
+
+    # 데이터셋 루트 경로 얻기 (프로젝트 id 기반)
+    
+    dataset_root_path = get_dataset_root_path(request.project_id)
+    
+    # 모델 로드
+    project_id = request.project_id
+    start_time = time.time()
+    send_slack_message(f"모델 로드 중 (projectId: {project_id})...", status="success")
+    model = get_model(project_id, request.m_key)
+    send_slack_message(f"모델 로드 완료 (projectId: {project_id}). 걸린 시간: {time.time() - start_time:.2f} 초", status="success")
+
+    # 이 값을 학습할때 넣으면 이 카테고리들이 학습됨
+    names = list(request.label_map)
+
+    # 레이블 변환기 (file_util.py/create_detection_train_label() 에 쓰임)
+    label_converter = {request.label_map[key]:idx for idx, key in enumerate(request.label_map)}
+    # key : 데이터에 저장된 프로젝트 카테고리 id
+    # value : 모델에 저장될 카테고리 id (모델에는 key의 idx 순서대로 저장될 것임)
+    
+    # 데이터 전처리: 학습할 디렉토리 & 데이터셋 설정 파일을 생성
+    start_time = time.time()
+    send_slack_message(f"데이터 전처리 시작: 학습 디렉토리 및 설정 파일 생성 중 (projectId: {project_id})...", status="success")
+    process_directories(dataset_root_path, names)
+    send_slack_message(f"데이터 전처리 완료 (projectId: {project_id}). 걸린 시간: {time.time() - start_time:.2f} 초",
+                       status="success")
+
+    # 데이터 전처리: 데이터를 학습데이터와 검증데이터로 분류
+    start_time = time.time()
+    send_slack_message(f"데이터 분류 중 (projectId: {project_id})...", status="success")
+    train_data, val_data = split_data(request.data, request.ratio)
+    send_slack_message(f"데이터 분류 완료 (projectId: {project_id}). 걸린 시간: {time.time() - start_time:.2f} 초",
+                       status="success")
+
+    # 데이터 전처리: 데이터 이미지 및 레이블 다운로드
+    start_time = time.time()
+    send_slack_message(f"데이터 다운로드 중 (projectId: {project_id})...", status="success")
+    download_data(train_data, val_data, dataset_root_path, label_converter)
+    send_slack_message(f"데이터 다운로드 완료 (projectId: {project_id}). 걸린 시간: {time.time() - start_time:.2f} 초",
+                       status="success")
+
+    # 학습 시작
+    start_time = time.time()
+    send_slack_message(f"학습 시작 (projectId: {project_id})...", status="success")
+    results = run_train(request, model, dataset_root_path)
+    send_slack_message(f"학습 완료 (projectId: {project_id}). 걸린 시간: {time.time() - start_time:.2f} 초", status="success")
+
+    # best 모델 저장
+    start_time = time.time()
+    send_slack_message(f"모델 저장 중 (projectId: {project_id})...", status="success")
+    model_key = save_model(project_id=project_id, path=join_path(dataset_root_path, "result", "weights", "best.pt"))
+    send_slack_message(f"모델 저장 완료 (projectId: {project_id}). 걸린 시간: {time.time() - start_time:.2f} 초", status="success")
+
+    result = results.results_dict
+
+    response = TrainResponse(
+        modelKey=model_key,
+        precision= result["metrics/precision(B)"],
+        recall= result["metrics/recall(B)"],
+        mAP50= result["metrics/mAP50(B)"],
+        mAP5095= result["metrics/mAP50-95(B)"],
+        accuracy=0,
+        fitness= result["fitness"]
+    )
+    send_slack_message(f"Detection train 성공 (projectId: {project_id}) {response}", status="success")
+
+    return response
+
+def split_data(data:list[TrainDataInfo], ratio:float):
+    try:
+        train_size = int(ratio * len(data))
+        random.shuffle(data)
+        train_data = data[:train_size]
+        val_data = data[train_size:]
+        
+        if not train_data or not val_data:
+            raise Exception("data size is too small")
+        return train_data, val_data
+    except Exception as e:
+        raise HTTPException(status_code=500, detail="exception in split_data(): " + str(e))
+
+def download_data(train_data:list[TrainDataInfo], val_data:list[TrainDataInfo], dataset_root_path:str, label_converter:dict[int, int]):
+    try:
+        for data in train_data:
+            process_image_and_label(data, dataset_root_path, "train", label_converter)
+
+        for data in val_data:
+            process_image_and_label(data, dataset_root_path, "val", label_converter)
+    except Exception as e:
+        raise HTTPException(status_code=500, detail="exception in download_data(): " + str(e))
+    
+def run_train(request, model, dataset_root_path):
+    try:
+        # 콜백 함수 정의
+        def send_data(trainer):
+            try:
+                # 첫번째 epoch는 스킵
+                if trainer.epoch == 0:
+                    return
+
+                # 남은 시간 계산(초)
+                left_epochs = trainer.epochs - trainer.epoch
+                left_seconds = left_epochs * trainer.epoch_time
+
+                # 로스 box_loss, cls_loss, dfl_loss
+                loss = trainer.label_loss_items(loss_items=trainer.loss_items)
+                data = ReportData(
+                    epoch=trainer.epoch,             # 현재 에포크
+                    total_epochs=trainer.epochs,     # 전체 에포크
+                    seg_loss=0,                      # seg_loss
+                    box_loss=loss["train/box_loss"], # box loss
+                    cls_loss=loss["train/cls_loss"], # cls loss
+                    dfl_loss=loss["train/dfl_loss"], # dfl loss
+                    fitness=trainer.fitness,         # 적합도
+                    epoch_time=trainer.epoch_time,   # 지난 에포크 걸린 시간 (에포크 시작 기준으로 결정)
+                    left_seconds=left_seconds        # 남은 시간(초)
+                )
+                # 데이터 전송
+                send_data_call_api(request.project_id, request.m_id, data)
+            except Exception as e:
+                # 예외 처리
+                print(f"Exception in send_data(): {e}")
+
+        # 콜백 등록
+        model.add_callback("on_train_epoch_start", send_data)
+
+        try:
+            # 비동기 함수로 학습 실행
+            results = model.train(
+                data=join_path(dataset_root_path, "dataset.yaml"),
+                name=join_path(dataset_root_path, "result"),
+                epochs=request.epochs,
+                batch=request.batch,
+                lr0=request.lr0,
+                lrf=request.lrf,
+                optimizer=request.optimizer,
+                patience=0
+            )
+        finally:
+            # 콜백 해제 및 자원 해제
+            model.reset_callbacks()
+            torch.cuda.empty_cache()
+        
+        # 마지막 에포크 전송
+        model.trainer.epoch += 1
+        send_data(model.trainer)
+        return results
+    except HTTPException as e:
+        raise e
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"exception in run_train(): {e}")
+
+@router.get("/memory")
+async def get_memory_status():
+    # GPU 메모리 정보 가져오기 (torch.cuda 사용)
+    if torch.cuda.is_available():
+        # 현재 활성화된 CUDA 디바이스 번호 확인
+        current_device = torch.cuda.current_device()
+
+        total_gpu_memory = torch.cuda.get_device_properties(current_device).total_memory
+        allocated_gpu_memory = torch.cuda.memory_allocated(current_device)
+        reserved_gpu_memory = torch.cuda.memory_reserved(current_device)
+
+        gpu_memory = {
+            "current_device" : current_device,
+            "total": total_gpu_memory / (1024 ** 3),  # 전체 GPU 메모리 (GB 단위)
+            "allocated": allocated_gpu_memory / (1024 ** 3),  # 현재 사용 중인 GPU 메모리 (GB 단위)
+            "reserved": reserved_gpu_memory / (1024 ** 3),  # 예약된 GPU 메모리 (GB 단위)
+            "free": (total_gpu_memory - reserved_gpu_memory) / (1024 ** 3)  # 사용 가능한 GPU 메모리 (GB 단위)
+        }
+        return gpu_memory
+    else:
+        raise HTTPException(status_code=404, detail="GPU가 사용 가능하지 않습니다.")

ai/app/api/yolo/model.py

@@ -0,0 +1,81 @@
+from fastapi import APIRouter, HTTPException, File, UploadFile
+from schemas.model_create_request import ModelCreateRequest
+from services.create_model import create_new_model, save_model
+from services.load_model import load_model
+from utils.file_utils import get_model_keys, delete_file, join_path, save_file, get_file_name
+import re
+from fastapi.responses import FileResponse
+
+router = APIRouter()
+
+@router.get("/info/projects/{project_id}/models/{model_key}", summary= "모델 관련 정보 반환")
+def get_model_info(project_id:int, model_key:str):
+    model_path = join_path("resources","projects", str(project_id), "models", model_key)
+    try:
+        model = load_model(model_path=model_path)
+    except FileNotFoundError:
+        raise HTTPException(status_code=404,
+                        detail= "모델을 찾을 수 없습니다.")
+    except Exception as e:
+        raise HTTPException(status_code=500, detail="model load exception: " + str(e))
+    # TODO: 학습치 등등 추가 예정
+    return {"type": model.task, "labelCategories":model.names}
+
+# project_id => model path 리스트 를 가져오는 함수
+@router.get("/projects/{project_id}", summary="project id 에 해당하는 모델 id 리스트")
+def get_model_list(project_id:int):
+    try:
+        return get_model_keys(project_id)
+    except FileNotFoundError:
+        raise HTTPException(status_code=404,
+                        detail= "프로젝트가 찾을 수 없거나 생성된 모델이 없습니다.")
+
+@router.post("/projects/{project_id}", status_code=201)
+def create_model(project_id: int, request: ModelCreateRequest):
+    model_key = create_new_model(project_id, request.project_type, request.pretrained)
+    return {"model_key": model_key}
+
+@router.delete("/projects/{project_id}/models/{model_key}", status_code=204)
+def delete_model(project_id:int, model_key:str):
+    model_path = join_path("resources", "projects", str(project_id), "models", model_key)
+    try:
+        delete_file(model_path)
+    except FileNotFoundError:
+        raise HTTPException(status_code=404,
+                        detail= "모델을 찾을 수 없습니다.")
+
+@router.post("/upload/projects/{project_id}")
+def upload_model(project_id:int, file: UploadFile = File(...)):
+    # 확장자 확인 확장자 새로 추가한다면 여기에 추가
+    if not file.filename.endswith(".pt"):
+        raise HTTPException(status_code=400, detail="Only .pt files are allowed.")
+
+    tmp_path = join_path("resources", "models", "tmp-"+file.filename)
+
+    # 임시로 파일 저장
+    try:
+        save_file(tmp_path, file)
+    except Exception as e:
+        raise HTTPException(status_code=500, detail="file save exception: "+str(e))
+    
+    # YOLO 모델 변환 및 저장
+    try:
+        model_path = save_model(project_id, tmp_path)
+        return {"model_path": model_path}
+    except Exception as e:
+        raise HTTPException(status_code=500, detail="file save exception: "+str(e))
+    finally:
+        # 임시파일 삭제
+        delete_file(tmp_path)
+
+
+@router.get("/download/projects/{project_id}/models/{model_key}")
+def download_model(project_id:int, model_key:str):
+    model_path = join_path("resources", "projects", str(project_id), "models", model_key)
+    try:
+        filename = get_file_name(model_path)
+        # 파일 응답 반환
+        return FileResponse(model_path, media_type='application/octet-stream', filename=filename)
+    except FileNotFoundError:
+        raise HTTPException(status_code=404,
+                        detail= "모델을 찾을 수 없습니다.")

ai/app/api/yolo/segmentation.py

@@ -0,0 +1,230 @@
+import time
+
+from fastapi import APIRouter, HTTPException
+from api.yolo.detection import get_classes, run_predictions, get_random_color, split_data, download_data
+from schemas.predict_request import PredictRequest
+from schemas.train_request import TrainRequest
+from schemas.predict_response import PredictResponse, LabelData
+from schemas.train_report_data import ReportData
+from schemas.train_response import TrainResponse
+from services.load_model import load_segmentation_model
+from services.create_model import save_model
+from utils.file_utils import get_dataset_root_path, process_directories, join_path
+from utils.slackMessage import send_slack_message
+from utils.api_utils import send_data_call_api
+
+router = APIRouter()
+
+@router.post("/predict")
+async def segmentation_predict(request: PredictRequest):
+    project_id = request.project_id
+    send_slack_message(f"Segmentation predict 요청 (projectId: {project_id}, 이미지 개수: {len(request.image_list)})",
+                       status="success")
+
+    # 모델 로드
+    start_time = time.time()
+    send_slack_message(f"모델 로드 중 (projectId: {project_id})...", status="success")
+    model = get_model(project_id, request.m_key)
+    send_slack_message(f"모델 로드 완료 (projectId: {project_id}). 걸린 시간: {time.time() - start_time:.2f} 초", status="success")
+
+    # 이미지 데이터 정리
+    start_time = time.time()
+    url_list = list(map(lambda x: x.image_url, request.image_list))
+    send_slack_message(f"이미지 데이터 정리 완료 (projectId: {project_id}). 걸린 시간: {time.time() - start_time:.2f} 초",
+                       status="success")
+
+    # 이 값을 모델에 입력하면 해당하는 클래스 id만 출력됨
+    classes = get_classes(request.label_map, model.names)
+
+    # 추론
+    start_time = time.time()
+    send_slack_message(f"Segmentation 추론 시작 (projectId: {project_id})...", status="success")
+    results = run_predictions(model, url_list, request, classes)
+    send_slack_message(f"Segmentation 추론 완료 (projectId: {project_id}). 걸린 시간: {time.time() - start_time:.2f} 초",
+                       status="success")
+
+    # 추론 결과 변환
+    start_time = time.time()
+    response = [process_prediction_result(result, image, request.label_map) for result, image in
+                zip(results, request.image_list)]
+    send_slack_message(f"Segmentation predict 성공 (projectId: {project_id}). 걸린 시간: {time.time() - start_time:.2f} 초",
+                       status="success")
+
+    return response
+
+# 모델 로드
+def get_model(project_id:int, model_key:str):
+    try:
+        return load_segmentation_model(project_id, model_key)
+    except Exception as e:
+        raise HTTPException(status_code=500, detail="load model exception: " + str(e))
+
+# 추론 결과 처리 함수
+def process_prediction_result(result, image, label_map):
+    try:
+        label_data = LabelData(
+            version="0.0.0",
+            task_type="seg",
+            shapes=[
+                {
+                    "label": summary['name'],
+                    "color": get_random_color(),
+                    "points": list(zip(summary['segments']['x'], summary['segments']['y'])),
+                    "group_id": label_map[summary['name']],
+                    "shape_type": "polygon",
+                    "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]
+        )
+    except Exception as e:
+        raise HTTPException(status_code=500, detail="model predict exception: " + str(e))
+
+    return PredictResponse(
+        image_id=image.image_id,
+        data=label_data.model_dump_json()
+    )
+
+
+@router.post("/train")
+async def segmentation_train(request: TrainRequest):
+    project_id = request.project_id
+
+    send_slack_message(f"Segmentation train 요청 (projectId: {project_id} 이미지 개수: {len(request.data)})", status="success")
+
+    # 데이터셋 루트 경로 얻기 (프로젝트 id 기반)
+    dataset_root_path = get_dataset_root_path(project_id)
+
+    # 모델 로드
+    start_time = time.time()
+    send_slack_message(f"모델 로드 중 (projectId: {project_id})...", status="success")
+    model = get_model(project_id, request.m_key)
+    send_slack_message(f"모델 로드 완료 (projectId: {project_id}). 걸린 시간: {time.time() - start_time:.2f} 초", status="success")
+
+    # 이 값을 학습할때 넣으면 이 카테고리들이 학습됨
+    names = list(request.label_map)
+
+    # 레이블 변환기
+    start_time = time.time()
+    label_converter = {request.label_map[key]: idx for idx, key in enumerate(request.label_map)}
+    send_slack_message(f"레이블 변환기 생성 완료 (projectId: {project_id}). 걸린 시간: {time.time() - start_time:.2f} 초",
+                       status="success")
+
+    # 데이터 전처리: 학습할 디렉토리 및 설정 파일 생성
+    start_time = time.time()
+    send_slack_message(f"데이터 전처리 중 (projectId: {project_id})...", status="success")
+    process_directories(dataset_root_path, names)
+    send_slack_message(f"데이터 전처리 완료 (projectId: {project_id}). 걸린 시간: {time.time() - start_time:.2f} 초",
+                       status="success")
+
+    # 데이터 전처리: 데이터를 학습데이터와 검증데이터로 분류
+    start_time = time.time()
+    train_data, val_data = split_data(request.data, request.ratio)
+    send_slack_message(f"데이터 분류 완료 (projectId: {project_id}). 걸린 시간: {time.time() - start_time:.2f} 초",
+                       status="success")
+
+    # 데이터 전처리: 데이터 이미지 및 레이블 다운로드
+    start_time = time.time()
+    send_slack_message(f"데이터 다운로드 중 (projectId: {project_id})...", status="success")
+    download_data(train_data, val_data, dataset_root_path, label_converter)
+    send_slack_message(f"데이터 다운로드 완료 (projectId: {project_id}). 걸린 시간: {time.time() - start_time:.2f} 초",
+                       status="success")
+
+    # 학습 시작
+    start_time = time.time()
+    send_slack_message(f"Segmentation 학습 시작 (projectId: {project_id})...", status="success")
+    results = run_train(request, model, dataset_root_path)
+    send_slack_message(f"Segmentation 학습 완료 (projectId: {project_id}). 걸린 시간: {time.time() - start_time:.2f} 초",
+                       status="success")
+
+    # best 모델 저장
+    start_time = time.time()
+    send_slack_message(f"모델 저장 중 (projectId: {project_id})...", status="success")
+    model_key = save_model(project_id=project_id, path=join_path(dataset_root_path, "result", "weights", "best.pt"))
+    send_slack_message(f"모델 저장 완료 (projectId: {project_id}). 걸린 시간: {time.time() - start_time:.2f} 초", status="success")
+
+    result = results.results_dict
+
+    response = TrainResponse(
+        modelKey=model_key,
+        precision=result["metrics/precision(M)"],
+        recall=result["metrics/recall(M)"],
+        mAP50=result["metrics/mAP50(M)"],
+        mAP5095=result["metrics/mAP50-95(M)"],
+        accuracy=0,
+        fitness=result["fitness"]
+    )
+    send_slack_message(f"Segmentation train 성공 (projectId: {project_id}) {response}", status="success")
+
+    return response
+    
+def run_train(request, model, dataset_root_path):
+    try:
+        # 데이터 전송 콜백함수
+        def send_data(trainer):
+            try:
+                # 첫번째 epoch는 스킵
+                if trainer.epoch == 0:
+                    return
+
+                # 남은 시간 계산(초)
+                left_epochs = trainer.epochs - trainer.epoch
+                left_seconds = left_epochs * trainer.epoch_time
+
+                # 로스 box_loss, cls_loss, dfl_loss
+                loss = trainer.label_loss_items(loss_items=trainer.loss_items)
+                data = ReportData(
+                    epoch=trainer.epoch,             # 현재 에포크
+                    total_epochs=trainer.epochs,     # 전체 에포크
+                    seg_loss=loss["train/seg_loss"], # seg_loss
+                    box_loss=0,                      # box loss          
+                    cls_loss=loss["train/cls_loss"], # cls loss
+                    dfl_loss=0,                      # dfl loss
+                    fitness=trainer.fitness,         # 적합도
+                    epoch_time=trainer.epoch_time,   # 지난 에포크 걸린 시간 (에포크 시작 기준으로 결정)
+                    left_seconds=left_seconds        # 남은 시간(초)
+                )
+                # 데이터 전송
+                send_data_call_api(request.project_id, request.m_id, data)
+            except Exception as e:
+                print(f"Exception in send_data(): {e}")
+
+        # 콜백 등록
+        model.add_callback("on_train_epoch_start", send_data)
+
+        try:
+            # 비동기 함수로 학습 실행
+            results = model.train(
+                data=join_path(dataset_root_path, "dataset.yaml"),
+                name=join_path(dataset_root_path, "result"),
+                epochs=request.epochs,
+                batch=request.batch,
+                lr0=request.lr0,
+                lrf=request.lrf,
+                optimizer=request.optimizer,
+                patience=0
+            )
+        finally:
+            # 콜백 해제 및 자원 해제
+            model.reset_callbacks()
+
+        # 마지막 에포크 전송
+        model.trainer.epoch += 1
+        send_data(model.trainer)
+        return results
+
+    except HTTPException as e:
+        raise e # HTTP 예외를 다시 발생
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"run_train exception: {e}")
+
+
+
+
+
+
+

ai/app/main.py

@@ -0,0 +1,53 @@
+from fastapi import FastAPI, Request
+from fastapi.exception_handlers import http_exception_handler, request_validation_exception_handler
+from fastapi.exceptions import RequestValidationError
+from starlette.exceptions import HTTPException
+from api.yolo.detection import router as yolo_detection_router
+from api.yolo.segmentation import router as yolo_segmentation_router
+from api.yolo.classfication import router as yolo_classification_router
+from api.yolo.model import router as yolo_model_router
+from utils.slackMessage import send_slack_message
+import time, torch, gc
+
+app = FastAPI()
+
+# 각 기능별 라우터를 애플리케이션에 등록
+app.include_router(yolo_detection_router, prefix="/api/detection", tags=["Detection"])
+app.include_router(yolo_segmentation_router, prefix="/api/segmentation", tags=["Segmentation"])
+app.include_router(yolo_classification_router, prefix="/api/classification", tags=["Classification"])
+app.include_router(yolo_model_router, prefix="/api/model", tags=["Model"])
+
+
+
+@app.middleware("http")
+async def resource_cleaner_middleware(request: Request, call_next):
+    start_time = time.time()
+    try:
+        response = await call_next(request)
+        return response
+    except Exception as exc:
+        raise exc
+    finally:
+        process_time = time.time() - start_time
+        if request.method != "GET":
+            send_slack_message(f"처리 시간: {process_time}초")
+        # gc.collect()
+        torch.cuda.empty_cache()
+    
+
+# 예외 처리기
+@app.exception_handler(HTTPException)
+async def custom_http_exception_handler(request:Request, exc):
+    body = await request.json()
+    send_slack_message(f"프로젝트 ID: {body['project_id']} - 실패! 에러: {str(exc)}", status="error")
+    return await http_exception_handler(request, exc)
+
+@app.exception_handler(RequestValidationError)
+async def validation_exception_handler(request:Request, exc):
+    send_slack_message(f"{request.url} - 요청 실패! 에러: {str(exc)}", status="error")
+    return await request_validation_exception_handler(request, exc)
+
+# # 애플리케이션 실행
+# if __name__ == "__main__":
+#     import uvicorn
+#     uvicorn.run("main:app", reload=True)

ai/app/schemas/model_create_request.py

@@ -0,0 +1,6 @@
+from pydantic import BaseModel
+from typing import Literal
+
+class ModelCreateRequest(BaseModel):
+    project_type: Literal["segmentation", "detection", "classification"]
+    pretrained:bool = True

ai/app/schemas/predict_request.py

@@ -0,0 +1,14 @@
+from pydantic import BaseModel, Field
+
+class ImageInfo(BaseModel):
+    image_id: int
+    image_url: str
+
+
+class PredictRequest(BaseModel):
+    project_id: int
+    m_key: str = Field("yolo8", alias="model_key")  # model_ 로 시작하는 변수를 BaseModel의 변수로 만들경우 Warning 떠서 m_key로 대체
+    label_map: dict[str, int] = Field(..., description="프로젝트 레이블 이름: 프로젝트 레이블 pk")
+    image_list: list[ImageInfo]  # 이미지 리스트
+    conf_threshold: float = Field(0.25, gt=0, lt= 1)
+    iou_threshold: float = Field(0.45, gt=0, lt= 1)

ai/app/schemas/predict_response.py

@@ -0,0 +1,23 @@
+from pydantic import BaseModel
+from typing import List, Optional, Tuple, Dict
+
+class Shape(BaseModel):
+    label: str
+    color: str
+    points: List[Tuple[float, float]]
+    group_id: Optional[int] = None
+    shape_type: str
+    flags: Dict[str, Optional[bool]] = {}
+
+class LabelData(BaseModel):
+    version: str
+    task_type: str
+    shapes: List[Shape]
+    split: str
+    imageHeight: int
+    imageWidth: int
+    imageDepth: int
+
+class PredictResponse(BaseModel):
+    image_id: int
+    data: str

ai/app/schemas/train_label_data.py

@@ -0,0 +1,28 @@
+from pydantic import BaseModel, Field
+
+
+class Segment(BaseModel):
+    x: float = Field(..., ge=0, le=1)
+    y: float = Field(..., ge=0, le=1)
+
+    def to_string(self) -> str:
+        return f"{self.x} {self.y}"
+
+class DetectionLabelData(BaseModel):
+    label_id: int = Field(..., ge=0)
+    center_x: float = Field(..., ge=0, le=1)
+    center_y: float = Field(..., ge=0, le=1)
+    width: float = Field(..., ge=0, le=1)
+    height: float = Field(..., ge=0, le=1)
+
+    def to_string(self) -> str:
+        return f"{self.label_id} {self.center_x} {self.center_y} {self.width} {self.height}"
+
+
+class SegmentationLabelData(BaseModel):
+    label_id: int
+    segments: list[Segment]
+
+    def to_string(self) -> str:
+        points_str = " ".join([segment.to_string() for segment in self.segments])
+        return f"{self.label_id} {points_str}"

ai/app/schemas/train_report_data.py

@@ -0,0 +1,12 @@
+from pydantic import BaseModel
+
+class ReportData(BaseModel):
+    epoch: int          # 현재 에포크
+    total_epochs: int   # 전체 에포크
+    seg_loss: float     # seg_loss
+    box_loss: float     # box loss
+    cls_loss: float     # cls loss
+    dfl_loss: float     # dfl loss
+    fitness: float      # 적합도
+    epoch_time: float   # 지난 에포크 걸린 시간 (에포크 시작 기준으로 결정)
+    left_seconds: float # 남은 시간(초)

ai/app/schemas/train_request.py

@@ -0,0 +1,22 @@
+from pydantic import BaseModel, Field
+from typing import Literal
+
+class TrainDataInfo(BaseModel):
+    image_url: str
+    data_url: str
+
+class TrainRequest(BaseModel):
+    project_id: int = Field(..., gt= 0)
+    m_key: str = Field("yolo8", alias="model_key")
+    m_id: int = Field(..., alias="model_id", gt= 0) # 학습 중 에포크 결과를 보낼때 model_id를 보냄
+    label_map: dict[str, int] = Field(..., description="프로젝트 레이블 이름: 프로젝트 레이블 pk")
+    data: list[TrainDataInfo]
+    ratio: float = Field(0.8, gt=0, lt=1) # 훈련/검증 분할 비율
+
+    # 학습 파라미터
+    epochs: int = Field(50, gt= 0, lt = 1000) # 훈련 반복 횟수
+    batch: int = Field(16, gt=0, le = 10000) # 훈련 batch 수[int] or GPU의 사용률 자동[float] default(-1): gpu의 60% 사용 유지
+    lr0: float = Field(0.01, gt= 0, lt= 1) # 초기 학습 가중치
+    lrf: float = Field(0.01, gt= 0, lt= 1) # lr0 기준으로 학습 가중치의 최종 수렴치 (ex lr0의 0.01배)
+    optimizer: Literal['auto', 'SGD', 'Adam', 'AdamW', 'NAdam', 'RAdam', 'RMSProp'] = 'auto'
+    

ai/app/schemas/train_response.py

@@ -0,0 +1,10 @@
+from pydantic import BaseModel
+
+class TrainResponse(BaseModel):
+    modelKey:str
+    precision: float
+    recall: float
+    mAP50: float
+    mAP5095: float
+    accuracy: float
+    fitness: float

ai/app/services/create_model.py

@@ -0,0 +1,51 @@
+from ultralytics import YOLO # Ultralytics YOLO 모델을 가져오기
+import os
+import uuid
+from services.load_model import load_model
+
+def create_new_model(project_id: int, type:str, pretrained:bool):
+    suffix = ""
+    type_list = {"segmentation": "seg", "classification": "cls"}
+    if type in type_list:
+        suffix = "-"+type_list[type]
+    # 학습된 기본 모델 로드
+    if pretrained:
+        suffix += ".pt"
+    else:
+        suffix += ".yaml"
+    model = YOLO(os.path.join("resources", "models" ,f"yolov8n{suffix}"))
+
+    # 모델을 저장할 폴더 경로
+    base_path = os.path.join("resources","projects",str(project_id),"models")
+    os.makedirs(base_path, exist_ok=True)
+
+    # 고유값 id 생성
+    unique_id = uuid.uuid4()
+    while os.path.exists(os.path.join(base_path, f"{unique_id}.pt")):
+        unique_id = uuid.uuid4()
+    model_path = os.path.join(base_path, f"{unique_id}.pt")
+
+    # 기본 모델 저장
+    model.save(filename=model_path)
+
+    return f"{unique_id}.pt"
+
+def save_model(project_id: int, path:str):
+    # 모델 불러오기
+    model = load_model(path)
+
+    # 모델을 저장할 폴더 경로
+    base_path = os.path.join("resources","projects",str(project_id),"models")
+    os.makedirs(base_path, exist_ok=True)
+
+    # 고유값 id 생성
+    unique_id = uuid.uuid4()
+    while os.path.exists(os.path.join(base_path, f"{unique_id}.pt")):
+        unique_id = uuid.uuid4()
+    model_path = os.path.join(base_path, f"{unique_id}.pt")
+
+    # 기본 모델 저장
+    model.save(filename=model_path)
+
+    return f"{unique_id}.pt"
+    

ai/app/services/load_model.py

@@ -0,0 +1,59 @@
+# ai_service.py
+
+from ultralytics import YOLO # Ultralytics YOLO 모델을 가져오기
+import os
+import torch
+
+def load_detection_model(project_id:int, model_key:str):
+    default_model_map = {"yolo8": os.path.join("resources","models","yolov8n.pt")}
+    # 디폴트 모델 확인
+    if model_key in default_model_map:
+        model = YOLO(default_model_map[model_key])
+    else:
+        model = load_model(model_path=os.path.join("resources", "projects",str(project_id),"models", model_key))
+
+    # Detection 모델인지 검증
+    if model.task != "detect":
+        raise TypeError(f"Invalid model type: {model.task}. Expected a DetectionModel.")
+    return model
+    
+def load_segmentation_model(project_id:int, model_key:str): 
+    default_model_map = {"yolo8": os.path.join("resources","models","yolov8n-seg.pt")}
+    # 디폴트 모델 확인
+    if model_key in default_model_map:
+        model = YOLO(default_model_map[model_key])
+    else:
+        model = load_model(model_path=os.path.join("resources", "projects",str(project_id),"models",model_key))
+        
+    # Segmentation 모델인지 검증
+    if model.task != "segment":
+        raise TypeError(f"Invalid model type: {model.task}. Expected a SegmentationModel.")
+    return model
+
+def load_classification_model(project_id:int, model_key:str):
+    default_model_map = {"yolo8": os.path.join("resources","models","yolov8n-cls.pt")}
+    # 디폴트 모델 확인
+    if model_key in default_model_map:
+        model = YOLO(default_model_map[model_key])
+    else:
+        model = load_model(model_path=os.path.join("resources", "projects",str(project_id),"models",model_key))
+        
+    # Segmentation 모델인지 검증
+    if model.task != "classify":
+        raise TypeError(f"Invalid model type: {model.task}. Expected a ClassificationModel.")
+    return model
+
+def load_model(model_path: str):
+    if not os.path.exists(model_path):
+        raise FileNotFoundError(f"Model file not found at path: {model_path}")
+    
+    try:
+        model = YOLO(model_path)
+        if (torch.cuda.is_available()):
+            model.to("cuda")
+            print("gpu 활성화")
+        else:
+            model.to("cpu")
+        return model
+    except:
+        raise Exception("YOLO model conversion failed: Unsupported architecture or invalid configuration.")

ai/app/utils/api_utils.py

@@ -0,0 +1,32 @@
+from schemas.train_report_data import ReportData
+from dotenv import load_dotenv
+import os, httpx
+
+
+def send_data_call_api(project_id:int, model_id:int, data:ReportData):
+    try:
+        load_dotenv() 
+        base_url = os.getenv("API_BASE_URL")
+        # main.py와 같은 디렉토리에 .env 파일 생성해서 따옴표 없이 아래 데이터를 입력
+        # API_BASE_URL = {url}
+        # API_KEY = {key}
+
+        # 하드코딩으로 대체
+        if not base_url:
+            base_url = "http://127.0.0.1:8080"
+        
+        headers = {
+            "Content-Type": "application/json"
+        }
+
+        response = httpx.request(
+            method="POST",
+            url=base_url+f"/api/projects/{project_id}/reports/models/{model_id}",
+            json=data.model_dump(),
+            headers=headers,
+            timeout=10
+        )
+        # status에 따라 예외 발생
+        response.raise_for_status()
+    except Exception as e:
+        print("report data failed: "+str(e))

ai/app/utils/file_utils.py

@@ -0,0 +1,161 @@
+import os
+import shutil
+import yaml
+from PIL import Image
+from schemas.train_request import TrainDataInfo
+from schemas.train_label_data import DetectionLabelData, SegmentationLabelData, Segment
+import urllib
+import json
+
+def get_dataset_root_path(project_id):
+    """데이터셋 루트 절대 경로 반환"""
+    return os.path.join(os.getcwd(), 'resources', 'projects', str(project_id), "train")
+
+def make_dir(path:str, init: bool):
+    """
+    path : 디렉토리 경로
+    init : 폴더를 초기화 할지 여부
+    """
+    if (os.path.exists(path) and init):
+        shutil.rmtree(path)
+    os.makedirs(path, exist_ok=True)
+
+def make_yml(path:str, model_categories):
+    data = {
+        "train": f"{path}/train",
+        "val": f"{path}/val",
+        "nc": len(model_categories),
+        "names": model_categories
+    }
+    with open(os.path.join(path, "dataset.yaml"), 'w') as f:
+        yaml.dump(data, f)
+
+def process_directories(dataset_root_path:str, model_categories:list[str]):
+    """학습을 위한 디렉토리 생성"""
+    make_dir(dataset_root_path, init=False)
+    make_dir(os.path.join(dataset_root_path, "train"), init=True)
+    make_dir(os.path.join(dataset_root_path, "val"), init=True)
+    if os.path.exists(os.path.join(dataset_root_path, "result")):
+        shutil.rmtree(os.path.join(dataset_root_path, "result"))
+    make_yml(dataset_root_path, model_categories)
+
+def process_image_and_label(data:TrainDataInfo, dataset_root_path:str, child_path:str, label_converter:dict[int,int]):
+    """이미지 저장 및 레이블 파일 생성"""
+    # 이미지 url로부터 파일명 분리
+    img_name = data.image_url.split('/')[-1]
+    img_path = os.path.join(dataset_root_path,child_path,img_name)
+
+    # url로부터 이미지 다운로드
+    urllib.request.urlretrieve(data.image_url, img_path)
+
+    # 파일명에서 확장자를 제거하여 img_title을 얻는다
+    img_title = os.path.splitext(os.path.basename(img_path))[0]
+
+    # 레이블 파일 경로
+    label_path = os.path.join(dataset_root_path, child_path, f"{img_title}.txt")
+
+    # 레이블 객체 불러오기
+    with urllib.request.urlopen(data.data_url) as response:
+        label = json.loads(response.read())
+
+    # 레이블 -> 학습용 레이블 데이터 파싱 후 생성
+    if label['task_type'] == "det":
+        create_detection_train_label(label, label_path, label_converter)
+    elif label["task_type"] == "seg":
+        create_segmentation_train_label(label, label_path, label_converter)
+
+def create_detection_train_label(label:dict, label_path:str, label_converter:dict[int, int]):
+    with open(label_path, "w") as train_label_txt:
+        for shape in label["shapes"]:
+            x1 = shape["points"][0][0]
+            y1 = shape["points"][0][1]
+            x2 = shape["points"][1][0]
+            y2 = shape["points"][1][1]
+            detection_label = DetectionLabelData(
+                label_id= label_converter[shape["group_id"]],     # 모델의 id  (converter : pjt category pk -> model category  id)
+                center_x= (x1 + x2) / 2 / label["imageWidth"],    # 중심 x 좌표
+                center_y= (y1 + y2) / 2 / label["imageHeight"],   # 중심 y 좌표
+                width= (x2 - x1) / label["imageWidth"],           # 너비
+                height= (y2 - y1) / label["imageHeight"]          # 높이
+            )
+            
+            train_label_txt.write(detection_label.to_string()+"\n") # str변환 후 txt에 쓰기
+
+def create_segmentation_train_label(label:dict, label_path:str, label_converter:dict[int, int]):
+    with open(label_path, "w") as train_label_txt:
+        for shape in label["shapes"]:
+            segmentation_label = SegmentationLabelData(
+                label_id = label_converter[shape["group_id"]],      # label Id
+                segments = [
+                    Segment(
+                        x=x / label["imageWidth"],          # shapes의 points 갯수만큼 x, y 반복
+                        y=y / label["imageHeight"]
+                    ) for x, y in shape["points"]
+                ]
+            )
+            train_label_txt.write(segmentation_label.to_string()+"\n")
+            
+def join_path(path, *paths):
+    """os.path.join()과 같은 기능, os import 하기 싫어서 만듦"""
+    return os.path.join(path, *paths)
+
+def get_model_keys(project_id:int):
+    path = os.path.join("resources","projects",str(project_id), "models")
+    if not os.path.exists(path):
+        raise FileNotFoundError()
+    files = os.listdir(path)
+    return files
+
+def delete_file(path):
+    if not os.path.exists(path):
+        raise FileNotFoundError()
+    os.remove(path)
+
+def save_file(path, file):
+    # 경로에서 디렉토리 부분만 추출 (파일명을 제외한 경로)
+    dir_path = os.path.dirname(path)
+    os.makedirs(dir_path, exist_ok=True)
+
+    with open(path, "wb") as buffer:
+        shutil.copyfileobj(file.file, buffer)
+
+def get_file_name(path):
+    if not os.path.exists(path):
+        raise FileNotFoundError()
+    return os.path.basename(path)
+
+def process_directories_in_cls(dataset_root_path:str, model_categories:list[str]):
+    """classification 학습을 위한 디렉토리 생성"""
+    make_dir(dataset_root_path, init=False)
+    for category in model_categories:
+        make_dir(os.path.join(dataset_root_path, "train", category), init=True)
+        make_dir(os.path.join(dataset_root_path, "test", category), init=True)
+    if os.path.exists(os.path.join(dataset_root_path, "result")):
+        shutil.rmtree(os.path.join(dataset_root_path, "result"))
+
+def process_image_and_label_in_cls(data:TrainDataInfo, dataset_root_path:str, child_path:str):
+    """이미지 저장 및 레이블 파일 생성"""
+    # 이미지 url로부터 파일명 분리
+    img_name = data.image_url.split('/')[-1]
+
+    # 레이블 객체 불러오기
+    with urllib.request.urlopen(data.data_url) as response:
+        label = json.loads(response.read())
+
+    if not label["shapes"]:
+        # assert label["shapes"], No Label. Failed Download" # AssertionError 발생
+        print("No Label. Failed Download")
+        return
+    label_name = label["shapes"][0]["label"]
+
+    label_path = os.path.join(dataset_root_path,child_path,label_name)
+
+    # url로부터 이미지 다운로드
+    if os.path.exists(label_path):
+        urllib.request.urlretrieve(data.image_url, os.path.join(label_path, img_name))
+    else:
+        # raise FileNotFoundError("No Label Category. Failed Download")
+        print("No Label Category. Failed Download")
+    # 레이블 데이터 중에서 프로젝트 카테고리에 해당되지않는 데이터가 있는 경우 처리 1. 에러 raise 2. 무시(+ warning)
+
+    

ai/app/utils/slackMessage.py

@@ -0,0 +1,27 @@
+import httpx
+import os
+
+SLACK_WEBHOOK_URL = "https://hooks.slack.com/services/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"
+
+def send_slack_message(message: str, status: str = "info"):
+    headers = {"Content-Type": "application/json"}
+
+    # 상태에 따라 다른 메시지 형식 적용 (성공, 에러)
+    if status == "error":
+        formatted_message = f":x: 에러 발생: {message}"
+    elif status == "success":
+        formatted_message = f":white_check_mark: {message}"
+    else:
+        formatted_message = message
+
+    # Slack에 전송할 페이로드
+    payload = {
+        "text": formatted_message
+    }
+
+    response = httpx.post(SLACK_WEBHOOK_URL, json=payload, headers=headers)
+
+    if response.status_code == 200:
+        return "Message sent successfully"
+    else:
+        return f"Failed to send message. Status code: {response.status_code}"

ai/app/utils/websocket_utils.py

@@ -0,0 +1,41 @@
+import websockets
+from websockets import WebSocketException
+
+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
+    
+class WebSocketConnectionException(WebSocketException):
+    def __init__(self, message="Failed to connect to WebSocket"):
+        super().__init__(message)

ai/environment.yml

@@ -0,0 +1,22 @@
+name: worlabel_ai_env
+channels:
+  - conda-forge
+  - pytorch
+  - nvidia
+  - defaults
+dependencies:
+  - python=3.10.10
+  - pytorch=2.3.1
+  - torchvision=0.18.1
+  - torchaudio=2.3.1
+  - pytorch-cuda=12.1
+  - fastapi
+  - uvicorn
+  - ultralytics
+  - dill
+  - boto3
+  - python-dotenv
+  - locust
+  - websockets
+  - httpx
+  - psutil

ai/locust/locustfile.py

@@ -0,0 +1,36 @@
+from locust import HttpUser, TaskSet, task, between
+
+class AIBehavior(TaskSet):
+    @task(weight = 1) # weight: 해당 task의 빈도수
+    def predict(self):
+        data = {
+            "project_id": 0,
+            "image_list": [
+                {
+                    "image_id": 12,
+                    "image_url": "test-data/images/image_000000001_jpg.rf.02ab6664294833e5f0e89130ecded0b8.jpg"
+                },
+                {
+                    "image_id": 23,
+                    "image_url": "test-data/images/image_000000002_jpg.rf.8270179e3cd29b97cf502622b381861e.jpg"
+                },
+                {
+                    "image_id": 47,
+                    "image_url": "test-data/images/image_000000003_jpg.rf.db8fd4730b031e35a60e0a60e17a0691.jpg"
+                }
+            ]
+        }
+        self.client.post("/api/detection", json=data)
+
+    # 앞으로 다른 API 또는 다른 data에 대해서 task 추가 가능
+
+class MyLocustUser(HttpUser):
+    wait_time = between(1,3)
+    tasks = [AIBehavior.predict]
+    host = "http://127.0.0.1:8000"
+
+# shell에 아래 명령어를 입력하여 실행(ai폴더 기준)
+# locust -f locust/locustfile.py
+# 또는
+# cd locust
+# locust

ai/requirements.txt

@@ -0,0 +1,12 @@
+fastapi==0.104.1
+uvicorn==0.30.6
+ultralytics==8.2.82
+ultralytics-thop==2.0.5
+
+--extra-index-url https://download.pytorch.org/whl/cu121
+torch==2.4.0+cu121
+--extra-index-url https://download.pytorch.org/whl/cu121
+torchaudio==2.4.0+cu121
+--extra-index-url https://download.pytorch.org/whl/cu121
+torchvision==0.19.0+cu121
+