[YOLOv8] custom training하여 바닥 Instance Segmentation하기

데이터

데이터 다운로드

데이터 다운로드

Download Dataset → Download dataset → Continue → Format을 COCO Segmentation으로 변경 → Continue

 

데이터 정보

TRAIN SET: 930개

VALID SET: 399개

데이터 형식: COCO

LABEL NAME: tehel

 

데이터 구조

data
	ㄴ train
		-annotations.json
		-이미지’s…
	ㄴ valid
		-annotations.json
		-이미지’s…

 

데이터 전처리

YOLOv8에서 COCO 데이터 형식을 지원한다고 들었으나,, 어떤 자료에서도 COCO형식 그대로 학습 하는걸 찾지 못했다. 따라서 COCO형식의 데이터를 YOLO형식의 데이터로 변환하여 학습했다. 데이터를 다운 받은 후 아래 파일을 실행하면 YOLO형식으로 변환된 데이터를 얻을 수 있다.

convertData.py(코드 설명)

1. 새로 생성할 데이터를 담을 폴더 생성
2. COCO DATA TYPE -> YOLO DATA TYPE
3. 이미지를 새로운 폴더로 이동
4. 1 인덱스 -> 0 인덱스(annotation file의 클래스)
5. 기존 데이터 폴더 삭제

변환된 데이터의 구조는 아래와 같다.

datasets
	ㄴ images
		ㄴ train
			- 이미지's...
		ㄴ val
			- 이미지's...
	ㄴ labels
		ㄴ train
			- 라벨's...
		ㄴ val
			- 라벨's...

 

학습 및 예측

data.yaml

학습된 데이터셋의 정보를 알려줄 data.yaml 작성한다.

data.yaml

path: [데이터셋의 경로]/datasets
train: images/train  # 훈련 이미지 경로
val: images/val   # 검증 이미지 및 경로

nc: 1 # 클래스의 개수
names: ["tehel"] #클래스의 이름

 

ultralytics와 YOLOV8

ultralytics를 통해 쉽게 yolov8을 학습 시킬 수 있으며 아래 명령어를 통해 다운받을 수 있다.

pip install ultralytics

 

yolov8의 Instance Segmentation에는 총 5가지의 모델이 존재하며 n에서 x모델로 갈수록 성능은 좋아지고 모델은 무거워진다.

이 문서에서는 yolov8n-seg.pt 모델을 기준으로 설명한다.

 

학습

아래 코드를 통해 모델을 불러오고 학습한다.

from ultralytics import YOLO

model = YOLO("yolov8n-seg.pt")  # Segmentation 모델
model.train(data="./data.yaml", epochs=50, imgsz=640)

 

예측

학습이 완료된 모델은 runs/segment/train/weights에 저장된다.

저장된 모델을 불러와 예측하는 코드는 아래와 같다.

from ultralytics import YOLO

model = YOLO("last.pt")  # Segmentation 모델
model.predict(source = "ASD.mp4", save=True, show_labels=True, show_conf=True, conf=0.5, save_txt=False, save_crop=False, line_width=2)

 

예측이 완료된 결과물

예측이 완료된 결과물은 runs/segment/predict에 저장된다.