Data Set 준비

필자는 roboflow를 이용하여 데이터 셋을 만들었고, YOLOv5 PyToch 형식으로 내보내기 하였다.

그러면 아래와 같이 다운로드 코드를 확인 할 수 있다.

Colab code

1. 설치

먼저 ultralytics를 설치하고 확인한다.

%pip install ultralytics

import ultralytics
ultralytics.checks()

2. Data Set 다운로드

roboflow를 설치하고 roboflow에서 복사한 Jupyter 다운로드 코드를 붙여넣는다.

%pip install roboflow

from roboflow import Roboflow
rf = Roboflow(api_key="api_key")
project = rf.workspace("my_workspace_name").project("project_name")
version = project.version(4)
dataset = version.download("yolov5")

api key와 workspace, project name은 공개하지 않기 위해 수정하였다.

3. Logger 선택

필자는 comet ML을 선택하였다.

logger = 'Comet' #@param ['Comet', 'TensorBoard']

if logger == 'Comet':
  %pip install -q comet_ml
  import comet_ml; comet_ml.init()
elif logger == 'TensorBoard':
  %load_ext tensorboard
  %tensorboard --logdir .

이때 Comet API key를 입력하라고 하는데 Comet ML 페이지에서 로그인 후 자신의 프로필을 누르면 API키를 아래와 같이 확인할 수 있다.

4. Train

mode를 train으로 설정하고 epochs는 100으로 설정하였다.

그리고 data는 자신의 data.yaml이 있는 곳의 경로로 설정해야 한다.

# Train YOLOv8s 100 epochs
!yolo task=detect mode=train model=yolov8s.pt data="your_path/data.yaml" epochs=100 imgsz=640

하지만 이렇게 수행하면 오류가 난다.

오류 내용을 확인해 보니 경로 문제인 것 같아서 data.yaml 파일을 확인해보았다.

확인해보니 train, val 경로가 잘못 작성 되어 있었다. 이 부분을 아래와 같이 수정해 주었다.

train: ../train/images
val: ../valid/images

필자는 이렇게 수정하였지만 각자 다를 수 있으니 경로를 확인해 보아야 한다.

5. Test

mode를 predict으로 설정하고 source부분에 자신이 테스트하고 싶은 이미지의 경로로 설정하면 된다.

# Test
!yolo task=detect mode=predict model=/content/runs/detect/train2/weights/best.pt conf=0.25 source=your_image_path

참고 자료

https://blog.roboflow.com/how-to-train-yolov8-on-a-custom-dataset/

https://colab.research.google.com/github/ultralytics/ultralytics/blob/main/examples/tutorial.ipynb#scrollTo=bpF9-vS_DAaf