TIL 2023.05.23

구글 colab 으로 YOLOv8 커스텀 트레이닝을 해보자

 

colab 준비:

구글 계정만 있으면 된다.

https://colab.research.google.com/

Google Colaboratory

Colab이란?

Colaboratory(줄여서 'Colab'이라고 함)을 통해 브라우저 내에서 Python 스크립트를 작성하고 실행할 수 있습니다.

• 구성이 필요하지 않음
• 무료로 GPU 사용
• 간편한 공유

학생이든, 데이터 과학자든, AI 연구원이든 Colab으로 업무를 더욱 간편하게 처리할 수 있습니다. Colab 소개 영상에서 자세한 내용을 확인하거나 아래에서 시작해 보세요.

라고 한다. colab에 대한 설명은 생략한다.

 

 

새 노트를 만들어 주고,

런타임 -> 런타임 유형 변경 -> 하드웨어 가속기: GPU 로 변경을 해주자.

 

욜로v8 관련 라이브러리 설치

!pip install ultralytics

 

학습하고자 하는 데이터 셋은 roboflow에서 찾아서 설치해주자.

https://universe.roboflow.com/

Roboflow Universe: Open Source Computer Vision Community

코드는 대략 이런 식

!curl -L "https://universe.roboflow.com/ds/25hLBsGCu1?key=???????" > roboflow.zip; unzip roboflow.zip; rm roboflow.zip

위 코드는 다운로드하고 압축도 풀어준다.

 

학습:

from ultralytics import YOLO
 
model = YOLO('yolov8n.pt') # 기본 욜로 모델로 전이 학습하기 위해 모델을 불러온다.
model.train(data='data.yaml', epochs=50)

따로 설정할 필요 없이 다운받은 데이터셋의 data.yaml만 불러오면 된다.

 

predict:

import cv2

model = YOLO("runs/detect/train/weights/best.pt") # 학습해서 나온 최선의 결과물이 저장된 경로
img = cv2.imread("your-image") # predict 하고자 하는 이미지를 불러오자
results = model.predict(source=img, save=True, save_txt=True)

결과는 run/detect/predict 폴더에 저장된다.