How to Train a Custom Ultralytics YOLOv8 Pose Estimation Model

FSA·2024년 10월 12일

keypoint_detection

목록 보기
4/19

사용자 정의 Ultralytics YOLOv8 자세 추정 모델 훈련 방법

  • 이 가이드에서는 Ultralytics YOLOv8 키포인트 탐지 모델을 사용자 정의 데이터셋에서 훈련하는 방법을 단계별로 설명합니다.
  • 우리는 접착제 스틱의 주요 지점을 식별하는 모델을 훈련하고, 이러한 지점을 사용해 이미지 내에서 접착제 스틱의 방향을 계산할 것입니다. 모델을 훈련하기 위해 다음 단계를 진행할 것입니다:
  1. 필수 종속성 설치
  2. Roboflow에서 키포인트 주석 추가
  3. 훈련에 사용할 데이터셋 생성
  4. 노트북에서 YOLOv8 키포인트 탐지 모델 훈련
  5. 객체의 방향을 계산하는 로직 작성
  • 이 튜토리얼을 따라 모델을 훈련할 수 있는 Colab 노트북을 만들었습니다.

1단계: 종속성 설치

pip install roboflow ultralytics

2단계: 데이터 라벨링

  • 키포인트 탐지 모델을 훈련하려면 이미지에 키포인트를 라벨링해야 합니다.
  • 이러한 키포인트는 이미지의 특정 지점에 대응됩니다.
  • 우리의 모델에서는 접착제 스틱의 이미지에 두 개의 주요 지점을 라벨링할 것입니다:
    • 캡과 스틱의 하단
  • 데이터를 라벨링하기 위해 우리는 Roboflow를 사용해 키포인트 데이터를 라벨링할 것입니다.
  • 이 가이드에서 사용하는 접착제 스틱 데이터셋은 Roboflow Universe에서 다운로드할 수 있습니다.
  • 무료 Roboflow 계정을 생성하세요. 계정을 생성한 후, Roboflow 대시보드로 가서 "Create New Project"(새 프로젝트 생성)을 클릭합니다.
  • 우리는 키포인트 탐지 모델을 만들고 싶습니다.
  • 이를 위해 팝업의 "Project type"(프로젝트 유형) 섹션에서 "See more"(더보기)을 클릭하고 "Keypoint detection"(키포인트 탐지)을 선택합니다.
  • 그런 다음 프로젝트 이름을 설정하고 프로젝트에 대한 라이센스를 선택합니다.
    • 데이터가 공개인 경우에만 라이센스를 설정해야 합니다.
    • 이는 Roboflow Universe에서 무료 플랜의 데이터가 공개된 경우에 해당합니다.
  • 프로젝트를 생성한 후 이미지를 업로드할 수 있습니다.
    • 이미지를 업로드하려면 이미지를 Roboflow 애플리케이션으로 드래그하세요.
    • 이미 COCO 키포인트 형식의 키포인트 주석이 있는 경우, 이를 이미지와 함께 업로드할 수 있습니다.
  • 이미지를 선택한 후 "Save and Continue"(저장 및 계속)을 클릭하여 Roboflow에 업로드합니다.
  • 다음으로, 식별하려는 각 클래스에 대한 스켈레톤을 설정해야 합니다.
    • 이는 Roboflow 클래스 관리 페이지에서 설정할 수 있습니다.
    • Roboflow 사이드바에서 "Classes"(클래스)를 클릭하세요. 그런 다음 모델 훈련에 사용할 모든 클래스를 추가하세요.
  • 이 가이드에서는 접착제 스틱 하나의 클래스를 식별하는 모델을 훈련할 것입니다.
  • 클래스를 추가한 후, 각 객체에 대해 라벨링할 지점을 정의해야 합니다.
  • 스켈레톤을 정의하려면 페이지의 그리드에서 아무 곳이나 클릭하여 포인트를 추가하세요.
  • 포인트를 클릭하면 포인트에 대한 이름을 지정할 수 있습니다.
  • 스켈레톤을 정의한 후 "Save"(저장)을 클릭합니다. 라벨링할 클래스마다 이 과정을 반복합니다. 이제 스켈레톤이 준비되었으므로 이미지를 주석 처리할 수 있습니다.

  • Roboflow 사이드바에서 "Annotate"(주석 추가)를 클릭하여 주석이 필요한 모든 이미지를 확인하세요. 이미지를 선택하면 Roboflow Annotate 인터페이스로 이동하여 이미지에 키포인트를 주석 처리할 수 있습니다.

  • 객체를 주석 처리하려면 오른쪽 사이드바에서 키포인트 주석 도구를 클릭하세요.
  • 그런 다음 주석을 추가할 객체 주위에 클릭하고 드래그하세요.
  • 클래스 이름을 선택하라는 메시지가 나타납니다.
  • 그런 다음 상자 안의 키포인트를 이미지의 각 지점 위치에 맞게 드래그할 수 있습니다.
  • 키포인트가 가려진 경우, 키포인트를 오른쪽 클릭한 후 "Mark as occluded"(가려진 것으로 표시)을 클릭하세요.
  • 이 가이드에서는 두 개의 지점을 주석 처리해야 합니다: 접착제 스틱의 상단(검은색 캡이 있는 부분)과 하단(스틱 하단의 절단선 부분)입니다.
  • 다음은 완전히 주석 처리된 이미지의 예시입니다:

3단계: 데이터셋 생성

  • Roboflow에서 이미지를 주석 처리한 후, 데이터셋을 생성할 수 있습니다.
  • 이 단계에서는 이미지에 전처리 단계와 증강을 적용할 수 있습니다.
  • Roboflow 대시보드의 왼쪽 사이드바에서 "Generate"(생성)을 클릭하여 데이터셋을 생성하세요.

  • 첫 번째 모델 버전을 훈련할 때는 Roboflow 대시보드에서 전처리 단계를 기본값으로 두는 것이 좋습니다.
  • 첫 번째 모델 버전에서는 증강을 추가하지 않는 것을 권장합니다.
  • 이는 나중에 증강을 사용한 실험 결과와 비교할 수 있도록 데이터에 대한 모델 성능을 이해하는 데 도움이 되기 때문입니다.

  • 증강을 추가하면 데이터셋에서 증강이 어떻게 보일지 미리볼 수 있습니다. 다음은 접착제 스틱 데이터셋에 적용된 밝기 증강의 미리보기입니다:

  • 데이터셋을 생성하려면 페이지 하단의 "Generate"(생성)를 클릭하세요.
    • 데이터셋이 생성됩니다.

  • 데이터셋이 생성되면 Roboflow에서 데이터셋을 사용할 수 있습니다.
  • Roboflow에서 모델을 훈련하거나, 데이터를 다운로드하여 자체 장치에서 YOLOv8 키포인트 탐지 모델을 훈련할 수 있습니다.

4단계: YOLOv8 키포인트 탐지 모델 훈련

  • pre-trained
    • 500 epoch을 학습시켰다.
    • A100에서 batch 48로 했는데, 기본값은 16이다.
  • 이제 데이터셋이 준비되었으므로, 키포인트 탐지 모델을 훈련할 수 있습니다.
  • 먼저 데이터셋을 다운로드해야 합니다. 다음 코드를 사용하여 데이터셋을 다운로드할 수 있습니다:
!pip install roboflow

from roboflow import Roboflow
rf = Roboflow(api_key="N7MTpQadXU3Tr0pjiHKZ")
project = rf.workspace("capjamesg").project("glue-keypoint-detection")
dataset = project.version(1).download("yolov5")

  • 이제 모델을 훈련할 수 있습니다.
  • 우리는 COCO 키포인트 체크포인트에서 모델을 훈련할 것입니다. 이는 대부분의 훈련 작업에 권장됩니다.
  • 모델을 훈련하려면 다음 코드를 추가하세요:
from ultralytics import YOLO
model = YOLO('yolov8n-pose.pt')

results = model.train(data=./'path/to/coco8-pose.yaml', epochs=100, imgsz=640)

  • 여기서 'path/to/coco8-pose.yaml’을 다운로드한 데이터셋의 YAML 파일 위치로 교체하세요.
train: ../train/images
val: ../valid/images
test: ../test/images

kpt_shape: [32, 3]
flip_idx: [24, 25, 26, 27, 28, 29, 22, 23, 21, 17, 18, 19, 20, 13, 14, 15, 16, 9, 10, 11, 12, 8, 6, 7, 0, 1, 2, 3, 4, 5, 31, 30]

nc: 1
names: ['pitch']

roboflow:
  workspace: roboflow-jvuqo
  project: football-field-detection-f07vi
  version: 15
  license: CC BY 4.0
  url: https://universe.roboflow.com/roboflow-jvuqo/football-field-detection-f07vi/dataset/15

  • 모델 훈련이 완료되면 모델을 로드하고 이미지에 추론을 실행할 수 있습니다:
from ultralytics import YOLO

model = YOLO('path/to/best.pt')

results = model(“image.png”)

print(results)

  • 다음은 접착제 스틱이 포함된 이미지에서 모델을 실행한 결과입니다. "results" 변수는 다음과 같습니다:
image 1/1 /content/datasets/Glue-Keypoint-Detection-1/test/images/IMG_4854_JPG.rf.db4d27601abf39708fcc2844695084d2.jpg: 640x640 1 glue, 50.3ms
Speed: 7.0ms preprocess, 50.3ms inference, 3.4ms postprocess per image at shape (1, 3, 640, 640)
ultralytics.engine.results.Keypoints object with attributes:

conf: tensor([[0.9489, 0.9854]], device='cuda:0')
data: tensor([[[230.2223, 167.3215,   0.9489],
         [126.4275, 192.6703,   0.9854]]], device='cuda:0')
has_visible: True
orig_shape: (640, 640)
shape: torch.Size([1, 2, 3])
xy: tensor([[[230.2223, 167.3215],
         [126.4275, 192.6703]]], device='cuda:0')
xyn: tensor([[[0.3597, 0.2614],
         [0.1975, 0.3010]]], device='cuda:0')
  • 다음 코드를 사용해 Google Colab에서 이미지에 키포인트를 그릴 수 있습니다:
from PIL import Image

r = results[0]
im_array = r.plot()  # 예측된 값을 BGR numpy 배열로 시각화
im = Image.fromarray(im_array[..., ::-1])  # RGB PIL 이미지
display(im)
  • 모델이 접착제 스틱의 위치를 성공적으로 식별했습니다.
profile
모든 의사 결정 과정을 지나칠 정도로 모두 기록하고, 나중에 스스로 피드백 하는 것

0개의 댓글