Yolo V8모델 커스텀 데이터 학습시키기
YOLO 모델로 커스텀 객체 인식을 가능하게 하기 위한 데이터셋 구축과 학습 방법 안내. labelImg를 통해 이미지 라벨링 후, 학습을 위한 데이터셋 폴더 구성 및 yaml 파일 설정 과정 설명. 추가로 YOLO 모델 학습에 필요한 파이썬 코드. 실시간 웹캠 객체 인식 코드로 최종 테스트
Yolo모델은 기본적으로 학습된 내용들이 제공된다. 하지만 기존의 학습된 사진말고 다른 객체를 인식시키고 싶은경우 커스텀 데이터를 학습시킬 수 있다.
먼저 학습시킬 데이터셋을 구축해야한다. 이미지들을 라벨링하는 과정은 필자가 쓴이전의 게시물을 참고하자.
참고로 이번 글은 gpu를 사용할 것이므로, cuda환경이 구축되어있다는 전제로 쓰여졌다.
Toby2718
위 게시물을 따라서 학습 데이터를 구축했다면 이제 학습을 시켜보자.
먼저 데이터 폴더 경로를 지정해주자.(아래처럼 폴더에 파일을 넣으면 된다)
dataset/
│
├── train/ # 학습 데이터 이미지 폴더
│ ├── image1.jpg
│ ├── image2.jpg
│ └── ...
│
├── val/ # 검증 데이터 이미지 폴더
│ ├── image1.jpg
│ ├── image2.jpg
│ └── ...
│
└── labels/ # 라벨 폴더
├── train/ # 학습 데이터의 라벨 (텍스트 파일들)
│ ├── image1.txt
│ ├── image2.txt
│ └── ...
│
└── val/ # 검증 데이터의 라벨 (텍스트 파일들)
├── image1.txt
├── image2.txt
└── ...
이제 데이터가 준비 되었으니 yaml파일을 생성해야한다.
아래 내용을 yaml파일을 폴더 내에 생성하자.
train: C:\Users\kkwjk2718\Desktop\test\train\imaes # 학습 이미지 경로
val: C:\Users\kkwjk2718\Desktop\test\valid\images # 검증 이미지 경로
test: C:\Users\kkwjk2718\Desktop\test\test\images # 테스트 이미지 경로
names: # 클래스 이름
0: card
1: red
2: box
위에는 데이터셋 경로, 아래는 클래스 이름을 지정해주면 된다
드디어 마지막이다
yolo모델 학습 파이썬 코드만 짜면 된다.
아래는 예시 코드이다.
from ultralytics import YOLO
def main():
model = YOLO('yolov8m.pt')
model.train(
data='yolo8Config.yaml',
epochs=50,
imgsz=640,
batch=10,
name='final',
device='cuda',
verbose=True
)
results = model.val()
print(results)
if __name__ == '__main__':
main()
Yolo모델 학습 실행 코드
여기서 model을 선택할 수 있는데 아래 사진처럼 성능별로 다양하게 존재하므로 각자 맞게 선택하면 된다. 또한 에포크와 배치사이즈도 적당하게 설정해주면 된다.
코드를 실행시키기 전 ultralytics라이브러리가 설치되어있지 않다면 pip로 설치해주자.
pip install ultralytics이제 실행시켜보면 학습이 잘 실행되고, gpu가 사용되고있는 모습을 볼 수 있다.
gpu가 살려달라고 소리지른다.
학습이 끝났다면 run폴더 안에서 평가, 학습관련 사진들과 weight파일을 받을 수 있다.
학습된 모델을 가지고 실시간으로 테스트해보자
아래 코드는 실시간으로 웹캠에서 이미지를 가져와서 객체인식하는 코드이다.
import cv2
from ultralytics import YOLO
model = YOLO('best.pt').to('cuda')
cap = cv2.VideoCapture(0)
if not cap.isOpened():
print("웹캠을 열 수 없습니다.")
exit()
while True:
ret, frame = cap.read()
if not ret:
print("프레임을 읽을 수 없습니다.")
break
results = model(frame)
annotated_frame = results[0].plot()
cv2.imshow('YOLOv8 Real-Time Detection', annotated_frame)
# 1ms 대기 후 ESC 키로 종료 (키보드 입력 확인)
if cv2.waitKey(1) & 0xFF == 27:
break
cap.release()
cv2.destroyAllWindows()
