Compare commits

...

133 Commits

Author SHA1 Message Date
187af772e4 Merge branch 'ai/develop' into 'master'
Ai/develop

See merge request s11-s-project/S11P21S002!312
2024-10-11 09:15:28 +09:00
김진현
b89e6ec827 Final: ai_final 2024-10-11 09:11:52 +09:00
김용수
b1c6282f52 Merge branch 'ai/fix/model' into 'ai/develop'
Fix: 모델 다운로드 path 오류 수정

See merge request s11-s-project/S11P21S002!292
2024-10-04 15:56:13 +09:00
김용수
dce2309f5b Merge branch 'ai/refactor/resource-cleanup' into 'ai/develop'
Feat: 모델 학습 patience(얼리스톱) 제거

See merge request s11-s-project/S11P21S002!291
2024-10-04 15:55:57 +09:00
김진현
be691ce04b Fix: 모델 다운로드 path 오류 수정 2024-10-04 15:53:34 +09:00
김진현
dae4ca4bf1 Merge branch 'ai/develop' of https://lab.ssafy.com/s11-s-project/S11P21S002 into ai/refactor/resource-cleanup 2024-10-04 15:43:05 +09:00
김진현
e428c6cf1f Feat: 모델 학습 patience(얼리스톱) 제거 2024-10-04 15:42:37 +09:00
김용수
4370631d5c Fix: 에러 해결 2024-10-04 15:16:47 +09:00
김진현
fbc2c1b46a Merge branch 'ai/refactor/train' into 'ai/develop'
Refactor: Segmentation 로깅 처리

See merge request s11-s-project/S11P21S002!289
2024-10-04 14:52:19 +09:00
김용수
4afe688f58 Merge branch 'ai/refactor/resource-cleanup' into 'ai/develop'
Feat: GET요청 시 Slack 메세지 제거

See merge request s11-s-project/S11P21S002!285
2024-10-04 14:47:40 +09:00
김용수
0871b1e951 Refactor: Segmentation 로깅 처리 2024-10-04 14:46:58 +09:00
김용수
4705fd318a Refactor: 슬랙 로깅 및 이미지 개수 처리 2024-10-04 14:31:21 +09:00
김진현
3fba3e71ef Feat: GET요청 시 Slack 메세지 제거 2024-10-04 14:24:47 +09:00
김진현
91fff1fb47 Merge branch 'ai/feat/memory' into 'ai/develop'
Feat: 메모리 확인

See merge request s11-s-project/S11P21S002!281
2024-10-04 13:30:16 +09:00
김용수
4a626ab050 Merge branch 'ai/refactor/resource-cleanup' into 'ai/develop'
Refactor: process_image_and_label url입출력 로직 수정

See merge request s11-s-project/S11P21S002!280
2024-10-04 13:03:07 +09:00
김용수
993b444c6e Feat: 메모리 확인 2024-10-04 13:02:33 +09:00
김진현
5d1bde161f Refactor: process_image_and_label url입출력 로직 수정 2024-10-04 12:49:17 +09:00
김용수
28c709e415 Feat: GPU 메모리 상태 응답 2024-10-04 11:36:21 +09:00
김용수
6fa912df5a Merge branch 'ai/refactor/resource-cleanup' into 'ai/develop'
Feat: 리소스 해제 관련 미들웨어 구현

See merge request s11-s-project/S11P21S002!273
2024-10-03 00:41:09 +09:00
김진현
6a9a8dd1ef Feat: 리소스 해제 관련 미들웨어 구현 2024-10-03 00:13:45 +09:00
김진현
7dabfaa467 Refactor: 비동기 run_in_threadpool() 모두 삭제 2024-10-03 00:01:54 +09:00
김용수
354676d867 Merge branch 'ai/refactor/exception-handle' into 'ai/develop'
Refactor: request 검증 classification 레이블 없을때 에러 처리

See merge request s11-s-project/S11P21S002!248
2024-09-30 15:39:30 +09:00
김진현
7962a2da29 Refactor: request 검증 classification 레이블 없을때 에러 처리 2024-09-30 15:24:44 +09:00
김진현
4dfa4b1c3e Refactor: train_label pydantic model 이용하여 데이터 유효성 검사 추가 2024-09-30 14:45:56 +09:00
김진현
1b7fea1415 Refactor: create_detection_train_label 레이블데이터의 이미지 바깥의 점에 대해 예외 처리 2024-09-30 12:12:44 +09:00
김진현
798f313e81 Feat: dotenv 이용하여 base_url 설정 2024-09-30 12:02:38 +09:00
정현조
4717ae9a0c Merge branch 'ai/feat/async' into 'ai/develop'
Feat: 오토레이블링, 학습 비동기 처리

See merge request s11-s-project/S11P21S002!230
2024-09-27 19:12:33 +09:00
김진현
7e75f41c64 Refactor: split_data() 예외처리 추가 2024-09-27 19:06:23 +09:00
김진현
7dd09182b8 Feat: 오토레이블링, 학습 비동기 처리 2024-09-27 19:03:18 +09:00
김진현
6bf4bba555 Merge branch 'ai/feat/classification' into 'ai/develop'
Fix: classification process_prediction_result() 에러 수정

See merge request s11-s-project/S11P21S002!227
2024-09-27 16:03:25 +09:00
김진현
7016d3a91e Fix: classification process_prediction_result() 에러 수정 2024-09-27 16:02:02 +09:00
김진현
df0a56565e Merge branch 'ai/feat/classification' into 'ai/develop'
Fix: segmentation get_model() parameter 수정

See merge request s11-s-project/S11P21S002!225
2024-09-27 15:50:48 +09:00
김진현
6a000be8d7 Fix: segmentation get_model() parameter 수정 2024-09-27 15:49:53 +09:00
김진현
8d89997369 Merge branch 'ai/feat/classification' into 'ai/develop'
Fix: classification result 수정

See merge request s11-s-project/S11P21S002!224
2024-09-27 15:46:52 +09:00
김진현
669562b712 Fix: classification result 수정 2024-09-27 15:46:02 +09:00
김태수
7ccd3b57d7 Merge branch 'ai/feat/classification' into 'ai/develop'
Fix: classification get_model() 파라미터 수정, 카테고리 없을 시 동작 수정

See merge request s11-s-project/S11P21S002!222
2024-09-27 15:17:52 +09:00
김진현
3affff85f0 Fix: classification get_model() 파라미터 수정, 카테고리 없을 시 동작 수정 2024-09-27 15:16:22 +09:00
김태수
9f269d8e5e Merge branch 'ai/feat/classification' into 'ai/develop'
Ai/feat/classification

See merge request s11-s-project/S11P21S002!220
2024-09-27 14:20:12 +09:00
김진현
4b6751a00b Refactor: Train response에 accuracy 추가 2024-09-27 14:18:59 +09:00
김진현
83737e566d Feat: classification 프로젝트 명세 변경에 따른 수정 2024-09-27 14:08:43 +09:00
김태수
15cc363e5c Merge branch 'ai/feat/segmentation' into 'ai/develop'
Fix: Detection, Segmentation 학습 레이블 컨버터 추가

See merge request s11-s-project/S11P21S002!216
2024-09-27 13:47:45 +09:00
김진현
02929ebc86 Fix: Detection, Segmentation 학습 레이블 컨버터 추가 2024-09-27 12:58:18 +09:00
정현조
ef021a9ca2 Merge branch 'ai/feat/segmentation' into 'ai/develop'
Feat: segmentation 프로젝트 명세 변경에 따른 수정

See merge request s11-s-project/S11P21S002!214
2024-09-27 12:33:45 +09:00
김진현
46c971494e Feat: segmentation 프로젝트 명세 변경에 따른 수정 2024-09-27 12:03:26 +09:00
김태수
2d6fdd23ff Merge branch 'ai/feat/detection' into 'ai/develop'
Fix: download_data paremeter 수정

See merge request s11-s-project/S11P21S002!213
2024-09-27 11:32:32 +09:00
김진현
0d4000b879 Fix: download_data paremeter 수정 2024-09-27 11:31:40 +09:00
김태수
959a16b59c Merge branch 'ai/feat/detection' into 'ai/develop'
Fix: get_model parameter 수정

See merge request s11-s-project/S11P21S002!212
2024-09-27 11:29:37 +09:00
김진현
8730b54d2d Fix: get_model parameter 수정 2024-09-27 11:28:52 +09:00
김용수
a93825ea8f Merge branch 'ai/feat/detection' into 'ai/develop'
Feat: detection/train 프로젝트 명세 변경에 따른 수정

See merge request s11-s-project/S11P21S002!211
2024-09-27 11:23:42 +09:00
김진현
0b38e54b21 Feat: detection/train 프로젝트 명세 변경에 따른 수정 2024-09-27 11:18:52 +09:00
김진현
2d8609077c Feat: detection/predict 프로젝트 명세 변경에 따른 수정 2024-09-27 10:38:01 +09:00
김진현
cae5fb5ae4 Feat: classification API 구현 2024-09-26 22:14:16 +09:00
김용수
7bcb4da904 Merge branch 'ai/feat/segmentation' into 'ai/develop'
Feat: 세그멘테이션 train response 수정

See merge request s11-s-project/S11P21S002!198
2024-09-26 21:26:10 +09:00
김진현
4d49b925dc Fix: 모델 생성 api에 오타 수정, project_id str로 변환 안한 버그 수정 2024-09-26 20:00:00 +09:00
김진현
d6b132c6ce Feat: 세그멘테이션 train response 수정 2024-09-26 19:53:56 +09:00
김용수
7f59d05a80 Merge branch 'ai/feat/segmentation' into 'ai/develop'
Fix: results_dict 오타 수정, raise 누락 수정

See merge request s11-s-project/S11P21S002!196
2024-09-26 19:29:09 +09:00
김진현
a3e0dc76dc Fix: results_dict 오타 수정, raise 누락 수정 2024-09-26 19:28:09 +09:00
김태수
57b54871f9 Merge branch 'ai/feat/segmentation' into 'ai/develop'
Refactor: train response 수정

See merge request s11-s-project/S11P21S002!193
2024-09-26 17:22:22 +09:00
김진현
336cc33384 Refactor: train response 수정 2024-09-26 17:20:37 +09:00
김태수
1cd5d77547 Merge branch 'ai/feat/segmentation' into 'ai/develop'
Feat: segmentation API 구현

See merge request s11-s-project/S11P21S002!187
2024-09-26 15:10:01 +09:00
김진현
6225c50914 Feat: segmentation API 구현 2024-09-26 15:04:31 +09:00
김용수
342795d7a3 Merge branch 'ai/refactor/detection' into 'ai/develop'
refactor: detection_train 리팩토링

See merge request s11-s-project/S11P21S002!184
2024-09-26 01:24:35 +09:00
김진현
6c9782a807 refactor: detection_train 리팩토링 2024-09-25 23:51:19 +09:00
김진현
42996f1417 Rename: send_data의 api의 함수명을 직관적으로 변경 2024-09-25 22:49:01 +09:00
김용수
a70ee0db39 Merge branch 'ai/fix/color' into 'ai/develop'
Fix: color # 수정

See merge request s11-s-project/S11P21S002!182
2024-09-25 22:05:28 +09:00
김진현
a1c0de1dfe Fix: color # 수정 2024-09-25 22:03:09 +09:00
김진현
86d76ab40a Merge branch 'ai/feat/error-handler' into 'ai/develop'
Feat: 에러핸들러 슬랙메세징 구현

See merge request s11-s-project/S11P21S002!181
2024-09-25 21:26:26 +09:00
김진현
7fcc0d9c28 Feat: 에러핸들러 슬랙메세징 구현 2024-09-25 21:24:30 +09:00
정현조
ac976afc84 Merge branch 'ai/refactor/predict' into 'ai/develop'
Feat: slack 알람설정, 학습 중 api 보내기 구현

See merge request s11-s-project/S11P21S002!180
2024-09-25 18:04:40 +09:00
김진현
82381a5292 Feat: slack 알람설정, 학습 중 api 보내기 구현 2024-09-25 17:59:02 +09:00
김용수
03d4a87907 Merge branch 'ai/refactor/predict' into 'ai/develop'
Refactor: 학습 레이블 data_url로 수정

See merge request s11-s-project/S11P21S002!175
2024-09-25 17:21:30 +09:00
김진현
7b10eeccaa Refactor: 학습 레이블 data_url로 수정 2024-09-25 16:06:45 +09:00
김태수
1187984698 Merge branch 'ai/refactor/predict' into 'ai/develop'
Refactor: predict, train 리팩토링

See merge request s11-s-project/S11P21S002!172
2024-09-25 15:28:21 +09:00
김진현
27cfd01fd7 Fix: 리팩토링에서 생긴 버그 수정 2024-09-25 15:24:14 +09:00
김진현
81b5145b41 Refactor: predict, train 리팩토링 2024-09-25 13:55:44 +09:00
김진현
20f8d9730a Merge branch 'ai/develop' of https://lab.ssafy.com/s11-s-project/S11P21S002 into ai/refactor/predict 2024-09-25 09:22:14 +09:00
김용수
70a62fe6f9 Feat: 슬랙 알림 메세지 연동: 2024-09-24 23:48:16 +09:00
김진현
cca395737a Merge branch 'ai/develop' of https://lab.ssafy.com/s11-s-project/S11P21S002 into ai/feat/229-data-process-for-train-stats 2024-09-24 09:16:01 +09:00
김진현
f8fae74e87 Merge branch 'ai/feat/backend' into 'ai/develop'
Feat: 백엔드 API 연동

See merge request s11-s-project/S11P21S002!143
2024-09-23 20:48:34 +09:00
김진현
22a8aa469a Merge branch 'ai/feat/backend' of https://lab.ssafy.com/s11-s-project/S11P21S002 into ai/feat/229-data-process-for-train-stats 2024-09-23 20:46:18 +09:00
김용수
a827a967c7 Feat: 백엔드 API 연동 2024-09-23 16:34:07 +09:00
정현조
dcd4799d20 Merge branch 'ai/feat/229-data-process-for-train-stats' into 'ai/develop'
Feat: 학습 통계 데이터 처리 구현 - S11P21S002-229

See merge request s11-s-project/S11P21S002!140
2024-09-23 16:28:19 +09:00
김진현
63b64e4e5f Feat: 학습 통계용 데이터 처리 구현 - S11P21S002-229 2024-09-23 16:24:56 +09:00
정현조
d9b6487a80 Merge branch 'ai/refactor/228-train-with-label-categories' into 'ai/develop'
Refactor: 레이블 카테고리를 포함한 학습 구현 및 리팩토링 - S11P21S002-228

See merge request s11-s-project/S11P21S002!128
2024-09-23 15:06:40 +09:00
김진현
1b17cfadad Refactor: 레이블 카테고리를 포함한 학습 구현 및 리팩토링 - S11P21S002-228 2024-09-23 13:06:48 +09:00
김진현
ae3506d2eb Rename: upload_tmp_model -> save_model 로 이름 변경 2024-09-23 10:45:20 +09:00
김용수
c6ef8ea5fc Merge branch 'ai/refactor/162-autolabel-with-label-categories' into 'ai/develop'
Refactor: 레이블 카테고리를 포함한 오토레이블링 구현 및 리팩토링

See merge request s11-s-project/S11P21S002!124
2024-09-23 10:23:09 +09:00
김진현
05e9a2c03b Refactor: 레이블 카테고리를 포함한 오토레이블링 구현 및 리팩토링 2024-09-23 09:46:42 +09:00
정현조
43f3d4104e Merge branch 'ai/refactor/model-path' into 'ai/develop'
Refactor: model_path 파라미터를 model_key로 변경

See merge request s11-s-project/S11P21S002!109
2024-09-20 17:53:46 +09:00
김진현
c7d07e5856 Merge branch 'ai/develop' of https://lab.ssafy.com/s11-s-project/S11P21S002 into ai/refactor/model-path 2024-09-20 13:56:22 +09:00
김진현
9b8d0927d7 Refactor: model_path 파라미터를 model_key로 변경 2024-09-20 13:54:46 +09:00
김용수
f9f3a0206d Merge branch 'ai/feat/52-ai-model-create' into 'ai/develop'
Feat: 추석동안 구현한 모델 관련 API 커밋 머지

See merge request s11-s-project/S11P21S002!97
2024-09-19 16:54:44 +09:00
김진현
9ce4c986d2 Feat: Detection prediction 웹소켓 연결 실패시에도 실행되도록 구현 2024-09-19 16:49:11 +09:00
김진현
47288e6bab Refactor: services/load_model 리팩토링 2024-09-19 16:48:19 +09:00
김진현
bcae5d5255 Refactor: 기본 모델 생성 구현 및 모델 생성 리팩토링 2024-09-19 01:15:41 +09:00
김진현
941df49f13 Feat: 모델 다운로드 API 구현 2024-09-18 16:27:14 +09:00
김진현
a98b3a0c9c Feat: 모델 업로드 API 구현 2024-09-18 16:09:18 +09:00
김진현
0503da4e41 Feat: 모델 삭제 API 구현 2024-09-18 02:20:27 +09:00
김진현
085aa71791 Feat: get_model_list API 구현 2024-09-18 01:45:56 +09:00
김진현
fc9c49b5f3 Feat: get_model_info API 구현 2024-09-18 01:08:40 +09:00
김진현
4e39dccd55 Feat: 모델 생성/학습된 모델 생성 API 구현 2024-09-18 00:13:01 +09:00
김진현
aae8faf11e Fix: rename 과정에서 생긴 import app.~에 의한 버그 수정
- 추가적으로 api endpoint 이름과 swagger 관련 설정
2024-09-18 00:02:02 +09:00
김진현
e2008e1d4d Docs: README.md에 resources 폴더 구조 설명 추가 2024-09-17 21:22:56 +09:00
김진현
a5ab6c716b Merge branch 'ai/develop' of https://lab.ssafy.com/s11-s-project/S11P21S002 into ai/feat/52-ai-model-create 2024-09-13 12:28:50 +09:00
김진현
1eacb7e90d Merge branch 'ai/feat/websocket' into 'ai/develop'
Feat: WebSocket 인스턴스 생성

See merge request s11-s-project/S11P21S002!65
2024-09-12 17:11:59 +09:00
김용수
fea960cc0d Feat: WebSocket 인스턴스 생성 2024-09-12 17:05:32 +09:00
김진현
68f6baa20a Rename: services 폴더 세분화 위해 ai_services.py -> load_model.py로 이름 변경 2024-09-12 10:23:55 +09:00
김용수
4fc13fa60c Feat: Fast API 소켓 열기 2024-09-12 00:53:45 +09:00
김용수
dbab504ab8 Merge branch 'ai/feat/165-locust-setting' into 'ai/develop'
Ai/feat/165 locust setting

See merge request s11-s-project/S11P21S002!58
2024-09-11 15:09:30 +09:00
김진현
3998e523da Merge branch 'ai/develop' of https://lab.ssafy.com/s11-s-project/S11P21S002 into ai/feat/165-locust-setting 2024-09-10 17:47:22 +09:00
김진현
af9096284c Feat: Locust 세팅 - S11P21S002-165 2024-09-10 17:46:12 +09:00
김진현
a5a4c49d62 Chore: Locust 패키지 추가 + torchvision 버전충돌 수정 2024-09-10 17:38:06 +09:00
김용수
10f60972b8 Merge branch 'ai/feat/117-detection-model-train' into 'ai/develop'
Feat: Detection 모델 학습 API 구현 - S11P21S002-117

See merge request s11-s-project/S11P21S002!56
2024-09-09 17:57:02 +09:00
김진현
75a12dd10c Feat: Detection 모델 학습 API 구현 - S11P21S002-117 2024-09-09 17:46:15 +09:00
김용수
1c74e539d9 Merge branch 'ai/feat/106-train-dataset-setting' into 'ai/develop'
Feat: Detection AI 모델 훈련 데이터 세팅 - S11P21S002-106

See merge request s11-s-project/S11P21S002!51
2024-09-09 15:04:28 +09:00
김진현
144d837361 Feat: Detection AI 모델 훈련 데이터 세팅 - S11P21S002-106 2024-09-08 16:46:59 +09:00
김진현
9abf99b992 Chore: environment.yml 패키지 수정 2024-09-05 17:07:25 +09:00
김태수
e86ce42166 Merge branch 'ai/feat/118-segmentation-autolabel' into 'ai/develop'
Feat: Segmentation 오토 레이블링 API 구현 - S11P21S002-118

See merge request s11-s-project/S11P21S002!48
2024-09-05 12:52:01 +09:00
김진현
1196476286 Feat: Segmentation 오토 레이블링 API 구현 - S11P21S002-118 2024-09-05 11:02:52 +09:00
정현조
4b56061f94 Merge branch 'ai/feat/detection-autolabel' into 'ai/develop'
Feat: Detection모델 로드 함수 보완

See merge request s11-s-project/S11P21S002!47
2024-09-05 10:45:00 +09:00
김진현
b0fc9f2de1 Merge branch 'ai/develop' of https://lab.ssafy.com/s11-s-project/S11P21S002 into ai/feat/118-segmentation-autolabel 2024-09-05 10:40:56 +09:00
김진현
924ad01962 Feat: Detection모델 로드 함수 보완
- Detection 모델인지 검증 부분 추가
- device 인자 및 gpu 이용 추가
2024-09-05 10:40:20 +09:00
김진현
8335f6d15d Fix: Detection 레이블 객체 파싱에서 imageDepth 수정 2024-09-05 10:36:17 +09:00
김태수
1f63944ac1 Merge branch 'ai/chore/env-setup' into 'ai/develop'
Chore: conda로 의존성 관리 추가

See merge request s11-s-project/S11P21S002!45
2024-09-05 09:49:23 +09:00
김진현
27e60ee68a Chore: conda로 의존성 관리 추가
- READ.me 및 노션 문서에 세팅 방법 적어둠
- conda로 관리하는 것이 Cuda 설치 등 더 유연함
- gpu 서버 또한 conda로 가상 환경 설정 예정
2024-09-04 22:26:22 +09:00
김진현
635c2a2033 Merge branch 'ai/feat/image' into 'ai/develop'
Refactor: API endPoint 변경 및 이미지 메모리 처리

See merge request s11-s-project/S11P21S002!43
2024-09-04 17:44:39 +09:00
김용수
6e87beab55 Refactor: API endPoint 변경 및 이미지 메모리 처리 2024-09-04 17:36:40 +09:00
홍창기
c7379d52ed Merge branch 'ai/feat/53-detection-autolabel' into 'ai/develop'
Feat: Detection 오토 레이블링 API 구현 - S11P21S002-53

See merge request s11-s-project/S11P21S002!37
2024-09-03 15:54:48 +09:00
김진현
84adb79c5e Refactor: 요청 응답 객체 구조 변경 2024-09-03 15:45:05 +09:00
김진현
8e77c225da Feat: Detection 오토레이블링 API 구현 - S11P21S002-53 2024-09-03 13:52:17 +09:00
김진현
03bba85028 Chore: .gitignore, requirements.txt 수정 2024-09-02 15:27:19 +09:00
김진현
1ed05d2de1 Doc: README.md 수정 2024-09-02 14:07:34 +09:00
김진현
4006ea1095 Init: AI 프로젝트 초기 설정 2024-09-02 13:50:43 +09:00
29 changed files with 1495 additions and 0 deletions

42
ai/.gitignore vendored Normal file
View File

@ -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

30
ai/README.md Normal file
View File

@ -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
- 훈련 데이터셋 저장

0
ai/app/__init__.py Normal file
View File

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

View File

View File

@ -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))

View File

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

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

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

View File

@ -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}")

53
ai/app/main.py Normal file
View File

@ -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)

View File

View File

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

View File

@ -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)

View File

@ -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

View File

@ -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}"

View File

@ -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 # 남은 시간(초)

View File

@ -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'

View File

@ -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

View File

View File

@ -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"

View File

@ -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.")

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

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

@ -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))

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

@ -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)

View File

@ -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}"

View File

@ -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)

22
ai/environment.yml Normal file
View File

@ -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

36
ai/locust/locustfile.py Normal file
View File

@ -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

12
ai/requirements.txt Normal file
View File

@ -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