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

Compare commits

base: jhynsoo:187af772e4b27fe9fa12f430e576084e919b1367
Branches Tags
jhynsoo:master
jhynsoo:ai/develop
jhynsoo:fe/develop
jhynsoo:be/develop
..
compare: jhynsoo:b51fd57e213cdfbe561d5b77f8fbe7b9f4b42964
Branches Tags
jhynsoo:master
jhynsoo:ai/develop
jhynsoo:fe/develop
jhynsoo:be/develop

No commits in common. "187af772e4b27fe9fa12f430e576084e919b1367" and "b51fd57e213cdfbe561d5b77f8fbe7b9f4b42964" have entirely different histories.
187af772e4 ... b51fd57e21

29 changed files with 0 additions and 1495 deletions
Show Stats Expand all files Collapse all files

42
ai/.gitignore vendored
View File

@ -1,42 +0,0 @@
# 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

30
ai/README.md
View File

@ -1,30 +0,0 @@
# 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
- 훈련 데이터셋 저장

0
ai/app/__init__.py
View File

0
ai/app/api/__init__.py
View File

0
ai/app/api/yolo/__init__.py
View File

199
ai/app/api/yolo/classfication.py
View File

@ -1,199 +0,0 @@
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))

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

@ -1,304 +0,0 @@
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가 사용 가능하지 않습니다.")

81
ai/app/api/yolo/model.py
View File

@ -1,81 +0,0 @@
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= "모델을 찾을 수 없습니다.")

230
ai/app/api/yolo/segmentation.py
View File

@ -1,230 +0,0 @@
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}")

53
ai/app/main.py
View File

@ -1,53 +0,0 @@
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)

0
ai/app/schemas/__init__.py
View File

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

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

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

@ -1,14 +0,0 @@
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)

23
ai/app/schemas/predict_response.py
View File

@ -1,23 +0,0 @@
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

28
ai/app/schemas/train_label_data.py
View File

@ -1,28 +0,0 @@
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}"

12
ai/app/schemas/train_report_data.py
View File

@ -1,12 +0,0 @@
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 # 남은 시간(초)

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

@ -1,22 +0,0 @@
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'

10
ai/app/schemas/train_response.py
View File

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

0
ai/app/services/__init__.py
View File

51
ai/app/services/create_model.py
View File

@ -1,51 +0,0 @@
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"

59
ai/app/services/load_model.py
View File

@ -1,59 +0,0 @@
# 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.")

0
ai/app/utils/__init__.py
View File

32
ai/app/utils/api_utils.py
View File

@ -1,32 +0,0 @@
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))

161
ai/app/utils/file_utils.py
View File

@ -1,161 +0,0 @@
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)

27
ai/app/utils/slackMessage.py
View File

@ -1,27 +0,0 @@
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}"

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

@ -1,41 +0,0 @@
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)

22
ai/environment.yml
View File

@ -1,22 +0,0 @@
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

36
ai/locust/locustfile.py
View File

@ -1,36 +0,0 @@
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

12
ai/requirements.txt
View File

@ -1,12 +0,0 @@
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