커스텀 데이터로 Yolo V8 Fine-Tuning

개요

ML 서비스 웹앱 배포 전주기 체험 프로젝트의 일환

 

데이터셋 수집, 정제, 모델학습 등 여러 작업을 번거롭게 하지 않아도 쉽게 가져와서 쓸수있는 뛰어난 모델들(메타의 SAM)이 많지만 전주기 체험을 목표로 하는만큼 잘 알려진 Yolo V8과 공개된 데이터셋을 활용해 Fine Tuning을 진행해보았다.

 

정리하는걸 좋아해서 쉽게 따라할 수 있도록 Colab 문서도 만들어보았다. (데이터셋 다운부터 학습까지)

yolo-v8-custom-dataset.ipynb

좌 입력 / 우 출력

---

모델 / 데이터셋 소개

모델 - Yolo V8 (Github)

Yolo 모델은 너무 유명해서 부가 설명이 필요없다.

- 가볍고 빠르면서도 적당한 성능을 뽑아줌

- Object Detection (특히 Realtime)에서 많이 사용되어 관련 자료가 풍부함

- Official Document 굉장히 상세함

버전이 올라가면서 라이선스 이슈가 있지만 위와 같은 이유로 사랑받는 모델.

데이터셋 - The Oxford-IIIT Pet Dataset

공개된 데이터 찾아보다 발견했는데 공개 데이터셋치고 품질도 좋고 장수도 많아서 사용하게 되었다.

- 유효하지 않은 이미지를 제외하고 보니 대략 7300장

- 각 이미지마다 annotation이 png 이미지 형태로 존재

- class, breed 등도 xml 데이터로 있다고 나와있는데 xml 파일은 3000장 정도만 있고 나머지는 없음

 

만드려는 서비스에 class나 breed 정보가 필요했다면 3000장으로만 학습하거나 어떻게든 구했을테지만 내가 만드려는 서비스에는 class 정보가 필요없기 때문에 그냥 xml 데이터를 사용하지 않기로 했다.

---

코드 리뷰

- 중요하지 않은 부분 생략 / 코랩 환경에서 작성

- 코랩 문서 참고하시면 이해가 쉽습니다.

 

annotation(png format) 확인

annotation 중 하나

직접 열어보면 그저 검은 이미지로 보인다.

annotation_path = f'dataset/annotations/trimaps/{images[sample_idx].split(".")[0]}.png'
sample_annotation = cv2.imread(annotation_path)

np.unique(sample_annotation)

# Output: array([1, 2, 3], dtype=uint8)

하지만 뜯어서 픽셀 값을 보면 0, 1, 2 세개의 값을 가진다.

각 값의 의미는 데이터셋 소개에도 나와있듯 아래와 같다.

1: Foreground - Detection 영역
2: Background - 배경
3: Not classified - etc

---

annotation image에서 유효한 영역 추출

위에서 로드한 sample_annotation에 아래와 같은 연산을 거쳐 binary mask로 만들 수 있다.

- 핵심은 로드한 이미지 픽셀값은 uint8 type으로 음수값을 갖지 않는다는 점

- np.sign은 부호를 추출하는 연산

sample_annotation = np.sign(sample_annotation_original - 2)
masked = np.multiply(sample_image, sample_annotation)

total = np.hstack((sample_image, sample_annotation*255, masked))

cv2_imshow(total)

# Output: Image

Output Image (좌: 원본 / 중: 연산 얻은 마스크 / 우: 마스크된 이미지)

마스크 이미지 좌측 위랑 우측 아래에 흰색 영역이 있는데 연산 결과가 아니라 원래 데이터가 그렇다.

---

annotation 변환

yolo에서는 주로 polygon 좌표가 담긴 txt format을 annotation으로 사용하기 때문에

png image에서 polygon을 추출해 txt 파일로 변환해야한다.

 

annotation_sample.txt 파일 예시

{class-index} {x1} {y1} {x2} {y2} ... {xn} {yn}

 

class, breed 등 정보가 담긴 xml 파일이 누락되어서 class-index는 0 하나로 통일

def annotationImage2yoloPolygonFormat(img: np.ndarray) -> List:
    img = np.sign(img - 2)

    if len(img.shape) == 3:
        img = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)

    H, W = img.shape

    contours, _ = cv2.findContours(img, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    contours_area = [cv2.contourArea(contour) for contour in contours]

    max_area = max(contours_area)
    max_index = contours_area.index(max_area)

    coords = []
    for point in contours[max_index]:
        x, y = point[0]
        coords.append(round(x / W, 3))
        coords.append(round(y / H, 3))

    return coords

위 함수는 아래 프로세스를 거친다.

1. numpy array로 로드된 annotation 이미지를 입력으로 받아

2. binary mask로 만든 뒤

3. cv2.findContours 메서드를 통해 폴리곤을 추출한다.

4. 이 때 여러개의 폴리곤이 추출될 수 있는데 그 중 가장 넓은 내부 영역을 가지는 폴리곤 하나만 선정하여

5. 1차원 List type으로 반환

sample_polygon = annotationImage2yoloPolygonFormat(sample_annotation_original)
print(f'polygon type: {type(sample_polygon)}')
print(f'polygon len: {len(sample_polygon)}')
print(f'polygon sample:{sample_polygon[:8]}')

# Output:
# polygon type: <class 'list'>
# polygon len: 756
# polygon sample:[0.143, 0.036, 0.14, 0.04, 0.13, 0.04, 0.127, 0.044]

---

변환된 annotation 복원해보기 (검증)

skimage 라이브러리의 polygon2mask 메서드를 사용해 변환한 폴리곤으로부터 마스크를 생성해본다.

정보 손실이 많이 없음을 알 수 있다.

from skimage.draw import polygon2mask

new_p = np.reshape(sample_polygon, (-1,2))
H, W, C = sample_annotation.shape

new_p = [[y*H, x*W] for x, y in new_p]

mask = polygon2mask((H, W), new_p)
mask = np.array(mask).astype(np.uint8)

cv2_imshow(sample_annotation * 255)
cv2_imshow(mask * 255)

# Output: Image

좌: 원본 / 우: 폴리곤으로부터 복원된 마스크

---

데이터 스플릿

모든 annotation 이미지에 대해 변환작업을 수행한 뒤 학습/검증 데이터로 데이터셋을 나눈다.

ultralytics에서 제공하는 util 사용

from ultralytics.data.utils import autosplit

autosplit(path='dataset/images', weights=(0.9, 0.1, 0.0))

# Output:
# Autosplitting images from dataset/images
# 100%|██████████| 7367/7367 [00:00<00:00, 30961.86it/s]

---

YAML 파일 작성

yolo에서는 하이퍼 파라미터, 데이터셋 정보 등의 메타데이터를 YAML 파일로 받는다.

하이퍼 파라미터는 yolo에서 제공하는 default 파일을 사용할거라 데이터 YAML만 쓴다.

%%writefile dataset/custom-data-seg.yaml
path: ../dataset
train: autosplit_train.txt
val: autosplit_val.txt

names:
  0: none
  
# Output:
# Writing dataset/custom-data-seg.yaml

경로와 클래스 정보를 기입해야한다.

클래스에 none 하나만 있는 이유는 해당 정보가 있어야 할 xml 파일 절반가량이 누락되어있어 사용하지 않았기 때문

---

학습

대망의 학습이지만 ultralytics에서 훌륭한 클래스와 메서드들을 제공하기에 import 제외 두 줄로 가능

 

- 여러 variation의 v8 중 nano 모델 로드한 뒤 pretrain weights 로드

- 위에서 써둔 custom-data-seg.yaml 로드하여 학습 진행

- 20 epoch, image size 640*640

from ultralytics import YOLO

model = YOLO('yolov8n-seg.yaml').load('yolov8n.pt')
model.train(data='dataset/custom-data-seg.yaml', epochs=20, imgsz=640)

# Output: ...

출력이 아주 길어 일부만 가져옴

Colab에서 무료로 제공하는 T4 GPU 달린 인스턴스를 사용해 20ep 학습에 약 1시간 걸렸다.

cls 정보는 넣지 않았으니 box_loss와 seg_loss, mAP 지표들만 살피면 된다.

경향을 봤을때 좀 더 학습 가능할 것 같지만 전주기 체험이라는 목적인만큼 모델을 정교하게 깎을 필요는 없기에 여기서 멈추는걸로

---

예측

예측 역시도 predict 메서드 한 줄로 할 수 있다.

result는 잘 포장된 predict 객체 형태이며 bounding box, binary mask 등 다양한 정보를 담고있다.

 

학습에 사용되지 않은 이미지를 로드, 예측한 뒤 그 중 mask 정보만 가져와서 씌워본다.

sample = cv2.imread('sample.jpg')

result = model.predict(sample, imgsz=640, conf=0.3, save=False, verbose=False)

h, w, _ = sample.shape

mask = result[0].masks.data.cpu().numpy()[0]
mask = cv2.resize(mask, dsize=(w, h))
mask = np.expand_dims(mask, axis=-1)

conc = np.hstack((sample, mask * sample))
cv2_imshow(conc)

# Output: Image

짜잔