YOLOv8 Segmentation모델을 Fine tune하기

이번 프로젝트에서 농구 경기 영상에서 경기장 영역을 segmentation하는 과정이 필요하다. 이 과정을 YOLOv8 seg모델을 fine tune하여 해결하려고 한다.

 

학습에 사용할 데이터셋은 다음과 같다. 해당 데이터셋을 사용하는 이유는, 단순한 경기장 영역을 라벨링한 것 뿐만아니라 3점라인 페인트 존 등 경기장 영역의 세부적인 요소까지 라벨링 되어있기 때문이다.

https://universe.roboflow.com/zy-vevvi/court-segmentation/dataset/4

court segmentation Instance Segmentation Dataset (v4, 2024-01-09 11:16pm) by Zy

 

YOLO학습을 위해서는 데이터셋과 데이터셋의 메타정보를 yolo모델에 전달해주는 data.yaml이 필요하다.

train: images/train
val: images/valid
test: images/test

# 클래스의 수와 클래스 이름 정의
names:
  0: "Center-circle"
  1: "Paint"
  2: "basketball-court"
  3: "three point line"

맨 위에 train, valid, test에 사용할 데이터셋 중 이미지들의 디렉토리 위치를 지정해주고 클래스의 정보를 정의해주면 된다. 위의 스크립트에서는 로컬디렉토리에서 images에 위치한 이미지 경로를 지정해주고있다.

 

Preprocessing

데이터과학에서 가장 중요한 작업은 전처리 과정이라고 할 수 있다. 이번 과정에서도 학습에 사용할 데이터셋을 전처리해줌으로써 yolo학습 데이터 템플릿에 맞추고, 모델 성능을 향상시키도록 한다.

데이터셋을 다운받으면 일반적인 yolo데이터셋의 형식과 다른것을 볼 수 있다. 데이터셋의 구조는 아래와 같아야 한다.

dataset/
├── images/
│   ├── train/
│   │   ├── train_image1.jpg
│   │   ├── train_image2.jpg
│   │   └── ...
│   ├── val/
│   │   ├── val_image1.jpg
│   │   ├── val_image2.jpg
│   │   └── ...
│   └── test/
│       ├── test_image1.jpg
│       ├── test_image2.jpg
│       └── ...
├── labels/
│   ├── train/
│   │   ├── train_image1.txt
│   │   ├── train_image2.txt
│   │   └── ...
│   ├── val/
│   │   ├── val_image1.txt
│   │   ├── val_image2.txt
│   │   └── ...
│   └── test/
│       ├── test_image1.txt
│       ├── test_image2.txt
│       └── ...

 

이를 맞춰주기 위해 폴더명을 데이터셋의 root폴더명(court segmentation.v4i.coco-mmdetection)을 images로 변경한다.

 

그 후 _annotations.coco.json안에 라벨링 되어있는 값들을 추출하여 labels폴더에 각 이미지에 대한 text파일을 작성해주어야 한다.

text file의 예시는 아래와 같다.

 

자세히 살펴보면 매 줄 첫번째는 라벨 id이고 그 이후는 segmentation polygon 좌표이다.

라벨 id는 0부터 시작하고

폴리곤 좌표는 x1 y1 x2 y2 x3 y3...으로 작성되어있다.

그리고 좌표가 모두 1미만의 값을 가지는것을 볼 수 있는데 이는 좌표를 정규화했기 때문이다.

정규화 공식
x : x좌표 / 이미지 너비
y : y좌표 / 이미지 높이

45 0.782016 0.986521 0.937078 0.874167 0.957297 0.782021 0.950562 0.739333 0.825844 0.561792 0.714609 0.420229 0.657297 0.391021 0.608422 0.4 0.0303438 0.750562 0.0016875 0.811229 0.003375 0.889896 0.0320156 0.986521
45 0.557859 0.143813 0.487078 0.0314583 0.859547 0.00897917 0.985953 0.130333 0.984266 0.184271 0.930344 0.386521 0.80225 0.480896 0.763484 0.485396 0.684266 0.39775 0.670781 0.3955 0.679219 0.310104 0.642141 0.253937 0.561234 0.155063 0.559547 0.137083
50 0.39 0.727063 0.418234 0.649417 0.455297 0.614125 0.476469 0.614125 0.51 0.590583 0.54 0.569417 0.575297 0.562354 0.601766 0.56 0.607062 0.536479 0.614125 0.522354 0.637063 0.501167 0.665297 0.48 0.69 0.477646 0.698828 0.494125 0.698828 0.534125 0.712938 0.529417 0.742938 0.548229 0.760594 0.564708 0.774703 0.550583 0.778234 0.536479 0.781766 0.531771 0.792359 0.541167 0.802937 0.555292 0.802937 0.569417 0.802937 0.576479 0.822359 0.576479 0.822359 0.597646 0.811766 0.607062 0.811766 0.618833 0.818828 0.637646 0.820594 0.656479 0.827641 0.687063 0.827641 0.703521 0.829406 0.727063 0.838234 0.708229 0.852359 0.729417 0.868234 0.750583 0.871766 0.792938 0.877063 0.821167 0.884125 0.861167 0.817062 0.92 0.734125 0.976479 0.711172 0.988229 0.48 0.988229 0.494125 0.967063 0.517062 0.912937 0.508234 0.832937 0.485297 0.788229 0.471172 0.774125 0.395297 0.729417
45 0.375219 0.0678333 0.375219 0.0590833 0.386828 0.0503542 0.424156 0.0315208 0.440797 0.0281458 0.464 0.0389167 0.525531 0.115583 0.611797 0.222521 0.676359 0.306583 0.678875 0.317354 0.677359 0.385271 0.66475 0.394687 0.588594 0.407458 0.417094 0.517771 0.280906 0.604521 0.0806562 0.722208 0.0256719 0.763917 0.00296875 0.809646 0 0.786104 0 0.745083 0 0.612583 0.03525 0.613271 0.0877187 0.626708 0.130594 0.626708 0.170437 0.6025 0.273844 0.548708 0.338906 0.507 0.509906 0.4115 0.604734 0.359042 0.596156 0.338188 0.595141 0.306583 0.595141 0.291792 0.579516 0.213104 0.516969 0.129042 0.498297 0.100792 0.466516 0.0987708 0.448875 0.0786042 0.405484 0.0705208 0.375219 0.0678333 0.28675 0.108375 0.282719 0.123167 0.267078 0.162854 0.266062 0.189083 0.245391 0.199833 0.203516 0.251625 0.187375 0.269771 0.159641 0.240188 0.101125 0.249604 0 0.287271 0 0.250271 0 0.245563 0.0975938 0.202521 0.203516 0.145354 0.251953 0.123167 0.28675 0.108375
49 0.587812 0.128229 0.612281 0.0965625 0.663391 0.0840833 0.690031 0.0908125 0.700109 0.10425 0.705859 0.133042 0.700109 0.143604 0.686422 0.146479 0.664828 0.153188 0.644672 0.157042 0.629563 0.175271 0.605797 0.181021 0.595 0.147437
49 0.7405 0.178417 0.733719 0.173896 0.727781 0.162583 0.729484 0.150167 0.738812 0.124146 0.747281 0.0981458 0.776109 0.0811875 0.804094 0.0845833 0.814266 0.102667 0.818516 0.115104 0.812578 0.133208 0.782906 0.151292 0.754063 0.172771
49 0.602656 0.178854 0.636125 0.167875 0.655172 0.165125 0.6665 0.162375 0.680391 0.155521 0.691719 0.153458 0.703047 0.154146 0.713859 0.162375 0.724156 0.174729 0.730844 0.193271 0.733422 0.217979 0.733938 0.244063 0.733422 0.281813 0.732391 0.295542 0.728266 0.300354 0.702016 0.294854 0.682969 0.28525 0.672156 0.270146
49 0.716891 0.0519583 0.683766 0.0103958 0.611688 0.0051875 0.568828 0.116875 0.590266 0.15325 0.590266 0.116875 0.613641 0.0857083 0.631172 0.0857083 0.6565 0.083125 0.679875 0.0883125 0.691563 0.0961042 0.711031 0.0649375

 

이 과정을 수행한 코드는 아래와 같다.

def convert_coco_to_yolo(coco_json_path, output_dir):
    with open(coco_json_path) as f:
        data = json.load(f)

    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    for image in data['images']:
        image_id = image['id']
        file_name = image['file_name']
        img_width = image['width']
        img_height = image['height']
        
        # YOLO 파일 이름 준비
        yolo_file_name = os.path.splitext(file_name)[0] + '.txt'
        yolo_file_path = os.path.join(output_dir, yolo_file_name)

        with open(yolo_file_path, 'w') as yolo_file:
            # 이미지에 해당하는 모든 어노테이션 찾기
            annotations = [ann for ann in data['annotations'] if ann['image_id'] == image_id]
            for ann in annotations:
                # COCO 포맷에서 YOLO 포맷으로 변환
                category_id = ann['category_id'] - 1  # YOLO는 클래스 인덱스가 0부터 시작
                segmentation = ann['segmentation'][0]
                normalized_segmentation = [coord / img_width if i % 2 == 0 else coord / img_height for i, coord in enumerate(segmentation)]
                
                segmentation_str = ' '.join(map(str, normalized_segmentation))

                yolo_file.write(f"{category_id} {segmentation_str}\n")

 

YOLOv8 Segmentation학습

위에서 전처리한 데이터셋을 학습시키는 yolo코드는 다음과 같다. gpu를 사용하며, 0번 gpu는 다른 모델 학습에 사용중이라 1번으로 지정하였다.

pre-trained모델인 yolov8n-seg.pt를 기반으로 fine tune하는 코드이다. epocs는 200, batch size는 16으로 지정하였다.

import torch
from ultralytics import YOLO

# 학습에 gpu사용
device = 'cuda:1' if torch.cuda.is_available() else 'cpu'
print(f'Using device: {device}')

# YOLOv8n-1280 모델 불러오기
model = YOLO('yolov8n-seg.pt').to(device)  # YOLOv8n 모델 사용

# 데이터셋 경로 설정 (YOLO 형식으로 구성된 데이터셋)
data_path = 'data.yaml'

# 학습 설정
epochs = 200
batch_size = 16
img_size = 640


# 모델 학습
results = model.train(
    data=data_path,
    epochs=epochs,
    batch=batch_size,
    imgsz=img_size,
    name='custom_yolov8_model'  # 저장될 모델의 이름 설정
)

 

결과

결과는 학습이 끝난 후 첨부하겠다.

---

트러블 슈팅

data.yaml에 따른 이미지 폴더 경로를 못 찾는 경우

Ultralytics폴더 안의 setting.yaml파일을 아래와같이 수정해주면 dataset, weights, run의 경로를 로컬 경로로 지정해줄 수 있다.

settings_version: 0.0.4
datasets_dir: dataset
weights_dir: weights
runs_dir: runs
...