Using Custom Datasets with YOLOv8 pretrained Detect model

오늘은 YOLOv8의 pretrained Detect model을 커스텀 데이터셋으로 돌려보는 작업을 따라해보려고 한다. 

https ://medium.com/zylapp/review-of-deep-learning-algorithms-for-object-detection-c1f3d437b852

 

 

YOLO의 경우 Object Detection 최초의 1-Stage Detector 방법론이다.

https://arxiv.org/abs/1506.02640

You Only Look Once: Unified, Real-Time Object Detection

 

자세하게 이해하려고 논문 리뷰하면 상당히 길기에, 나동빈님유튜브나 스탠포드cs231n, 고대 연구실 유튜브를 보고 개념정리 후에 논문을 읽어보고 필자의 순서대로 사용해보는 것을 추천한다. 정말 좋은 시대에 태어났다...

 

 

예전에 나온 방법론이지만, 여러 최적화가 이루어지면서 최근에는 YOLOv8까지 나오게 됐다. 

성능도 우수하고 FPS도 상당히 빠르기 때문에, 최근에도 Object Detection을 활용할 수 있는 분야에서 자주 사용된다.

하지만 YOLO는 요새 모델에 대한 논문 발표는 하지 않고있다.

 

필자는 아래 유튜브 영상의 방법을 따라가며 진행했다.

https://youtu.be/gRAyOPjQ9_s?si=y-8OsLi-y_d1j9Li

 

---

필자의 환경

윈도우 11, RTX3060 laptop

(리눅스 환경에서도 무난히 돌아갈 것으로 보인다)

아나콘다, VScode

CUDA 11.3

가상환경설정 

아나콘다로 가상환경 구축한다. (역시 conda)

conda create -n yolov8 python=3.9

 

데이터 수집

이 부분에서 본인이 가지고 있는 커스템 데이터셋이 있다면 생략해도 된다.

# simple_image_download 이미지 데이터 다운로드용 라이브러리 설치
pip install simple_image_download==0.4

 

simple_image_download 라이브러리를 이용해서 “building workers”와 관련된 이미지 데이터 다운로드를 하고자 한다.

 

필자는 바탕화면에 yolov8이라는 폴더를 생성했다.

나중에 아나콘다에서 cd 경로를 이용해서 가상환경을 폴더안으로 들어오게 해주자.

 

yolov8폴더 안에  download_images .py을 만든다.

텍스트 파일을 아무거나 생성해서 확장자 파일만 바꿔주면된다.

 

ex) 이름없음.txt -> (확장자명까지 다 지우고)download_images.py 

이름없음.txt -> (확장자명까지 다 지우고)download_images.py

 

그 후에 .py 파일을 VScode와 같은 툴로 아래와 같이 작성해준다.

from simple_image_download import simple_image_download as simp

response = simp.simple_image_download

keywords = ["building workers"]

for kw in keywords:
    response().download(kw, 200)

 

저장 잊지 말기

 

그 후에 .py파일을 실행시킨다.

python download_images.py

 

그러면 아래와 같이 자동적으로 폴더가 생성되고 폴더 안에 이미지가 다운로드 된다. 

 

 

그리고 yolov8 폴더에 images라는 폴더를 생성하고,

생성된 이미지 파일 중 쓸모 없는 것들을 제거하며 파일을 옮긴다.

(필자는 직접 라벨링할것이기 때문에, 너무 많으면 힘들어서 데이터를 몇 개 삭제했다.)

 

그 후, 필요 없어진 폴더를 삭제한다.

images 폴더 생성후 이미지 파일 옮기기

 

우리가 가지고 있는 이미지 데이터들은 label(정답)데이터가 존재하지 않는다.

즉, 직접 라벨링 과정을 거쳐서 정답 레이블들을 만들어 줘야한다.

 

데이터 수집

라벨링은 사실 가장 귀찮은 작업이며, 손이 많이간다. 

본인이 가지고 있는 데이터에 라벨 데이터가 있다면, 이 작업을 수행할 필요가 없다.

대신 yolo가 원하는 라벨 데이터에 맞춰서 변경해야하는 작업이 수행되어야 할 것이다.

(코드 짜서 변경하면 된다. 변경하는 작업이 직접 라벨링하는 것보다 훨씬 쉬울 것이다.)

 

라벨링 하는 방법은 여러가지가 있지만, 일단 영상대로 직접해보도록하자.

먼저 라벨링하는 labelImg 라이브러리를 설치한다.

pip install labelImg

 

설치완료 됐다면, 아래 코드를 입력하자.

labelImg

 

 

그러면 아래와 같은 창이 뜨게 된다.

 

open폴더를 지정해주고 File -> Open Dir

 

change Save Dir을 지정해준다.  File -> Change Save Dir

 

Create RectBox를 클릭해서 정답이 될 부분에 바운딩 박스를 쳐준다.

그 후 본인이 박스쳐준 곳에 해당하는 classes를설정해준다.

 

※ Create RectBox위에 PascalVOC를 눌러서 YOLO로 바꿔줘야 yolo에 맞는 라벨 데이터가 생성된다.

왼쪽 아래 Create RectBox를 클릭

classes 지정

YOLO로 해줘야함!!

 

사진마다 하나하나 라벨링을 해줬다면 아래 사진과 같이 생성이 됐을 것이다.

마지막에 classes.txt 를 확인해보면 설정한 classes들이 무엇인지 확인할 수 있다.

라벨링 데이터 확인

 

데이터 분리

Train / Val 폴더를 만들고 데이터를 분리해보자.

 

val 폴더 안에도 images와 labels 폴더를 만들고 train과 비율을 맞춰서 몇 개 파일을 선택하여 넣어준다.

참고로 확장자 명 외의 파일 명은 동일해야한다.

 

val 폴더에 안넣은 이미지와 라벨데이터는 전부 train에 넣어준다.

그냥 yolov8폴더에 있는 폴더 2개(images, labels)를 train에 넣어주면된다.

이 과정에서 classes.txt는 yolov8폴더로 빼준다.

classes.txt는 yolov8폴더로 빼준다.

 

yolov8폴더에 있는 폴더 2개(images, labels)를 train에 넣어주면된다.

 

그러면 아래와 같은 폴더 상황이 된다.

 

마지막으로, 모델을 돌릴 때 필요한 데이터 카테고리 확인 파일인 data_custom.yaml을 생성한다.

.py파일과 동일하게 아무 텍스트 파일을 만들고 확장자 명을 바꿔서 생성해준다.

 

그 후, 아래 내용을 적어서 저장한다.

# 경로는 알아서 바꿔주세요 
train: C:\Users\jeauk\OneDrive\바탕 화면\yolov8\train 
# 경로는 알아서 바꿔주세요 
val: C:\Users\jeauk\OneDrive\바탕 화면\yolov8\val

nc: 2

names: ["jacket","hat"]

경로를 다음과 같이 했지만 에러가걸리기 때문에 전체 경로를 작성해준다.

 

yolo 라이브러리 설치 (with PyTorch)

yolo를 본격적으로 사용하기 위해서 설치해보자

https://github.com/ultralytics/ultralytics

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

 

설치는 매우 간단하다.

pip install ultralytics

 

단, 설치 후 torch 버전은 cpu버전으로 다운로드 된다.

따라서 우리는 torch gpu버전으로 깔아줘야한다.

본인이 가지고 있는 cuda에 맞춰서 가장 최신으로 pytorch를 다운로드 해준다.

--upgrade 붙여주는게 포인트다.

 

필자는https://pytorch.org/get-started/previous-versions/ 여기서 맞는 버전을 다운로드 해줬다.

https://pytorch.org/get-started/locally/

Start Locally

https://wookdata.tistory.com/62

윈도우 파이토치 설치 방법[How to install Pytorch in Windows 11](CUDA 환경 구성)

 

# --upgrade 붙여주기!!
pip install --upgrade torch==1.12.1+cu113 torchvision==0.13.1+cu113 torchaudio==0.12.1 --extra-index-url https://download.pytorch.org/whl/cu113

 

 

아래 코드를 돌렸을때, true가 나오면 성공!

import torch
torch.cuda.is_available()

 

YOLO train 시작

우리는 YOLOv8의 pretrained Detect model을 돌릴것이다.

https://docs.ultralytics.com/tasks/detect/#models

Detect

 

모델 종류가 다양하게 있지만 YOLOv8m으로 돌려보자.

 

아래 코드를 실행시킨다. 만약 코드가 돌아가지 않는다면 GPU의 램 문제이니 배치사이즈를 조정해준다.

yolo task=detect mode=train epochs=100 data=data_custom.yaml model=yolov8m.pt imgsz=640

 

다행히 문제없이 돌아갔다.

 

61 epochs이 돌아가고 학습이 종료 되었다. 더이상 학습률이 올라가지 않는다고 자동적으로 판단하고 학습을 중지한 것이다. (이것도 코드 수정으로 가능하긴하다.) 

 

runs\detect\train에 훈련된 결과가 저장되어있다고 한다.

폴더로 가보면,

여러가지 시각화에 도움이되는 metric 결과와 loss, 가중치 파일 ,

그리고 결과 metric들이 csv파일로도 저장되어있다.

상당히 편한 모델임을 확인할 수 있다!!!!

Test

마지막으로 다온 훈련 결과를 Test해보자.

대충 train이든 val이든 아무 이미지를 yolov8 폴더로 꺼내서 코드를 돌린다.

yolo task=detect mode=predict model=best.pt show=True conf=0.5 source=test.jpeg

conf=0.5: 신뢰도 임계값을 0.5로 설정.

 

객체 감지에서 모델은 각 감지에 신뢰도 점수를 할당하여 감지된 객체가 올바르게 식별되었는지 여부를 나타낸다. 이를 설정하면 conf=0.5신뢰도 점수가 50% 이상인 탐지만 고려하도록 모델에 지시하게 된다.

 

모델이 확실하지 않은 탐지를 필터링하여 잠재적으로 오탐지 수를 줄이는 방법이다.

 

전반적으로 hat은 학습이 잘된편이지만, jacket의 경우 조금 처참한 학습을 보였다.

(데이터의 양, 잘못된 라벨링의 문제) 

 

이건 좀 웃겨서 가져왔다. 전등을 안전모로 detect한다. 

 

결과는 솔직히 좀 아쉽지만, 내가 라벨링을 이상하게 했다는 점(나름 꼼꼼히 했는데 ㅠ),

데이터의 양이 매~~~우 적었다는 점에서 나름 성공적인 튜토리얼이었다라고 할 수 있겠다. 

 

라벨링이 귀찮은 사람들을 위한... 구린 나의 라벨링 데이터를 뿌리도록 하겠다...

https://github.com/wook2jjang/YOLOv8_Custom_Dataset

GitHub - wook2jjang/YOLOv8_Custom_Dataset: Using Custom Datasets with YOLOv8

결론

결론적으로,  YOLOv8의 pretrained Detect model을 이용해서,

아주아주 쉽게 본인이 가진 custom 데이터셋을 transfer learning 시키고 돌려볼 수 있는 좋은 기회라 생각한다!

 

다음에는 라벨 데이터가 있는 다량의 데이터를 학습시켜 돌려보도록 하겠다.