diff --git a/ai/.gitignore b/ai/.gitignore new file mode 100644 index 0000000..b7baa62 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/ai/README.md b/ai/README.md new file mode 100644 index 0000000..44be447 --- /dev/null +++ b/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 +- 훈련 데이터셋 저장 \ No newline at end of file diff --git a/ai/app/__init__.py b/ai/app/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/ai/app/api/__init__.py b/ai/app/api/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/ai/app/api/yolo/__init__.py b/ai/app/api/yolo/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/ai/app/api/yolo/classfication.py b/ai/app/api/yolo/classfication.py new file mode 100644 index 0000000..3de0eea --- /dev/null +++ b/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)) + + diff --git a/ai/app/api/yolo/detection.py b/ai/app/api/yolo/detection.py new file mode 100644 index 0000000..2fcf0dc --- /dev/null +++ b/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가 사용 가능하지 않습니다.") \ No newline at end of file diff --git a/ai/app/api/yolo/model.py b/ai/app/api/yolo/model.py new file mode 100644 index 0000000..9a50831 --- /dev/null +++ b/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= "모델을 찾을 수 없습니다.") diff --git a/ai/app/api/yolo/segmentation.py b/ai/app/api/yolo/segmentation.py new file mode 100644 index 0000000..c372cd0 --- /dev/null +++ b/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}") + + + + + + + diff --git a/ai/app/main.py b/ai/app/main.py new file mode 100644 index 0000000..b119930 --- /dev/null +++ b/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) diff --git a/ai/app/schemas/__init__.py b/ai/app/schemas/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/ai/app/schemas/model_create_request.py b/ai/app/schemas/model_create_request.py new file mode 100644 index 0000000..0f17676 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/ai/app/schemas/predict_request.py b/ai/app/schemas/predict_request.py new file mode 100644 index 0000000..e4cf3b3 --- /dev/null +++ b/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) diff --git a/ai/app/schemas/predict_response.py b/ai/app/schemas/predict_response.py new file mode 100644 index 0000000..3570b32 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/ai/app/schemas/train_label_data.py b/ai/app/schemas/train_label_data.py new file mode 100644 index 0000000..372d3f7 --- /dev/null +++ b/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}" \ No newline at end of file diff --git a/ai/app/schemas/train_report_data.py b/ai/app/schemas/train_report_data.py new file mode 100644 index 0000000..1e4a89f --- /dev/null +++ b/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 # 남은 시간(초) \ No newline at end of file diff --git a/ai/app/schemas/train_request.py b/ai/app/schemas/train_request.py new file mode 100644 index 0000000..a06899f --- /dev/null +++ b/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' + diff --git a/ai/app/schemas/train_response.py b/ai/app/schemas/train_response.py new file mode 100644 index 0000000..222b3ce --- /dev/null +++ b/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 \ No newline at end of file diff --git a/ai/app/services/__init__.py b/ai/app/services/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/ai/app/services/create_model.py b/ai/app/services/create_model.py new file mode 100644 index 0000000..b5c6b8e --- /dev/null +++ b/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" + \ No newline at end of file diff --git a/ai/app/services/load_model.py b/ai/app/services/load_model.py new file mode 100644 index 0000000..ad12f97 --- /dev/null +++ b/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.") diff --git a/ai/app/utils/__init__.py b/ai/app/utils/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/ai/app/utils/api_utils.py b/ai/app/utils/api_utils.py new file mode 100644 index 0000000..be0d809 --- /dev/null +++ b/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)) diff --git a/ai/app/utils/file_utils.py b/ai/app/utils/file_utils.py new file mode 100644 index 0000000..b14b85b --- /dev/null +++ b/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) + + \ No newline at end of file diff --git a/ai/app/utils/slackMessage.py b/ai/app/utils/slackMessage.py new file mode 100644 index 0000000..633b01c --- /dev/null +++ b/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}" diff --git a/ai/app/utils/websocket_utils.py b/ai/app/utils/websocket_utils.py new file mode 100644 index 0000000..02b5474 --- /dev/null +++ b/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) \ No newline at end of file diff --git a/ai/environment.yml b/ai/environment.yml new file mode 100644 index 0000000..2161aa9 --- /dev/null +++ b/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 diff --git a/ai/locust/locustfile.py b/ai/locust/locustfile.py new file mode 100644 index 0000000..9772b65 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/ai/requirements.txt b/ai/requirements.txt new file mode 100644 index 0000000..a2cb5ad --- /dev/null +++ b/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 +