Yolov8 - 내 커스텀 데이터를 이용해 학습하기 / 훈련 데이터와 검증 데이터로 나누기

 

사용하기에 앞서 https://github.com/ultralytics/ultralytics

위 링크에 접속하여 Clone한 후에 내 로컬에서 열어줍니다. 

 

언어는 Python을 이용해 작업을 진행해보겠습니다.

GitHub - ultralytics/ultralytics: NEW - YOLOv8 🚀 in PyTorch > ONNX > CoreML > TFLite

위 모델을 이용하기 위해서는 반드시 .yaml 파일이 필요합니다. 

 

아래는 제가 작업한 yaml파일이며 train: 부분과 val: 부분에 각각의 경로를 설정해주시면 됩니다.

 

# Ultralytics YOLO 🚀, GPL-3.0 license


train: C:\Users\your_train_path
val: C:\Users\your_val_path

nc: 5 # Number_of_Class

# Classes
names: ["NG_Hole","Dent","Scratch","Wrinkle","foreign_substance"]

 

파일이 보통 train과 val로 나뉘어져 있지 않을 때, 즉 훈련 데이터과 검증 데이터로 나누는 코드는 아래와 같습니다.

import os
import glob

path = 'D:/your_path' #파일이 위치한 경로

# 모든 이미지 파일을 리스트에 넣는다


all_files = []
for i in range(1, 7):  # 디렉토리 이름이 1부터 6까지의 번호로 지정되어 있음
    jpg_files = glob.glob(f"{path}/{i}/*.jpg")  # 각 디렉토리의 모든 jpg 파일을 찾음
    for jpg in jpg_files: # 각 jpg 파일에 대해
        txt_file = os.path.splitext(jpg)[0] + '.txt' # 해당하는 txt 파일의 이름을 만듭니다.
        if os.path.exists(txt_file): # txt 파일이 실제로 존재하는지 확인합니다.
            all_files.append(jpg) # 해당하는 txt 파일이 있다면, jpg 파일을 리스트에 추가합니다.

import random
# 모든 파일을 랜덤하게 섞습니다.
random.shuffle(all_files)

# 훈련 데이터와 검증 데이터로 나눈다.
split_ratio = 0.8  # 80% for training, 20% for validation
split_index = int(split_ratio * len(all_files))

train_files = all_files[:split_index]
val_files = all_files[split_index:]

# train.txt와 val.txt 파일에 각각 훈련 데이터와 검증 데이터의 파일 경로를 기록합니다.
with open("train.txt", "w") as f:
    f.write("\n".join(train_files))

with open("val.txt", "w") as f:
    f.write("\n".join(val_files))

import shutil

# 기본 디렉토리 설정
root_dir = 'D:/your_dir'

# 훈련 및 검증에 필요한 디렉토리를 생성
train_image_dir = os.path.join(root_dir, 'mobis/train/images/')
train_label_dir = os.path.join(root_dir, 'mobis/train/labels/')
val_image_dir = os.path.join(root_dir, 'mobis/val/images/')
val_label_dir = os.path.join(root_dir, 'mobis/val/labels/')

# 디렉토리가 없으면 생성, 이미 있으면 그냥 두기
os.makedirs(train_image_dir, exist_ok=True)
os.makedirs(train_label_dir, exist_ok=True)
os.makedirs(val_image_dir, exist_ok=True)
os.makedirs(val_label_dir, exist_ok=True)


# 훈련용 파일 이동
for file in train_files:
    shutil.copy(file, train_image_dir)  # 이미지 파일을 새 디렉토리로 복사
    txt_file = os.path.splitext(file)[0] + '.txt'  # 해당 이미지와 연관된 txt 파일 이름을 얻기
    shutil.copy(txt_file, train_label_dir)  # 라벨 파일을 새 디렉토리로 복사
    
    
    
 # 검증용 파일 이동
for file in val_files:
    shutil.copy(file, val_image_dir)  # 이미지 파일을 새 디렉토리로 복사
    txt_file = os.path.splitext(file)[0] + '.txt'  # 해당 이미지와 연관된 txt 파일 이름을 얻기
    shutil.copy(txt_file, val_label_dir)  # 라벨 파일을 새 디렉토리로 복사

import ultralytics
ultralytics.checks()
from ultralytics import YOLO
model = YOLO('yolov8n.pt') # 사용할 모델 불러오기

print(type(model.names),len(model.names)) # 불러온 모델 확인
print(model.names) # 불러온 모델 확인

import os
os.environ['KMP_DUPLICATE_LIB_OK']='TRUE'

model.train( data='mobis.yaml', epochs=50, patience=30, batch=1 ) # 여러 매개변수가 있음, 오픈소스 제공해주는 깃허브에서 확인 가능

 

 

처음에는 batch= 파라미터를 64, 32로 진행을 해봤는데, 배치 크기가 너무 커서 GPU 메모리를 초과해서 학습이 진행되지 않는 문제가 발생 됐다, 위 코드는 테스트용이기에 1로 조절하고 하니 잘 되는 것을 확인 이후에 최적의 파라미터들을 찾아서 다시 학습을 진행을 해봐야겠다.