파이참으로 YOLOv5 모델 학습 with Custom dataset

 Custom dataset을 활용하여 pycharm으로 YOLO v5 모델을 통한 프로젝트를 진행하는 내용에 대해 작성해보겠습니다. 이번 프로젝트에서는 일반적인 캔(음료수캔, 커피캔, 맥주캔)과 부탄캔 이미지를 분류하는 작업을 수행하였습니다.

 

1. 이미지 데이터 구하기

일반캔 이미지 데이터셋

부탄캔 이미지 데이터셋

 프로젝트를 진행하기 위해 먼저 일반적인 캔(이하 일반캔)과 부탄캔 이미지를 수집하였습니다. 일반캔의 경우 한국지능정보사회진흥원에서 운영하는 AI허브의 생활 폐기물 이미지 중 캔류에서 원천 데이터를 얻어내었고, 부탄캔 이미지의 경우 직접 촬영하거나 인터넷에 있는 이미지 데이터를 수집해서 활용하였습니다.

 

AI-Hub

 

2. 데이터 라벨링하기

 YOLO 모델에서 학습을 진행시키기 위해서는 데이터 라벨링 과정이 선행되어야 합니다. 이번 프로젝트에서는 makesense 라는 도구를 활용하여 손쉽게 라벨링 과정을 수행하였습니다. Makesense를 통해 라벨링을 수행하기 위해서는 아래 링크에서 Get Started를 누르고 라벨링을 수행할 이미지들을 추가한 뒤 Object Detection을 통해 다음 화면으로 넘어간 뒤 Create labels 창에서 라벨링에 필요한 수량 만큼 + 버튼을 누른 뒤 원하는 이름을 넣고 Start project를 선택하면 빠르게 진행할 수 있습니다.

 

Make Sense

 

 이번 프로젝트에서는 10,000장의 일반캔 및 부탄캔 이미지에 대해 라벨링 작업을 약 3~4주 가량에 걸쳐 수행하였으며, 사각형이 물체를 정확하게 표시하는 것을 중점적으로 진행하였습니다.

makesense를 통한 라벨링 작업

 해당 이미지들의 라벨링이 완료되면 상단 메뉴 중 Actions에서 Export Annotations 버튼을 누른 뒤, YOLO 학습을 위해 하단 이미지의 첫번째 YOLO format을 체크하고 Export를 통해 결과물을 저장합니다.

makesense

 다운로드 받은 결과물의 합축을 해제하면 아래 이미지와 같은 텍스트 파일이 생성되고 해당 파일의 맨 앞 숫자는 class를 나머지 숫자들은 라벨링한 사각형의 좌표를 의미하게 됩니다.

라벨링 된 결과

3. train / val / test로 데이터셋 나누기

 라벨링 작업이 모두 마무리 되면 데이터셋을 세부류로 나누게 됩니다. 이번 프로젝트에서는 총 10,000장의 데이터를 train : val : test = 8 : 1 : 1 비율로 나누어 진행하였습니다.

 Custom_dataset이라는 폴더를 만들어 주고 그 하위 폴더로 test, train, val을 각각 만들어 둔 뒤, 각각의 하위폴더에는 images, labels이라는 폴더를 만들고 shuffing 후 파일을 이동할 이미지 갯수에 맞추어 8,000개, 1,000개, 1,000개를 넣어주었습니다.

 

Shuffing과 이미지 이동을 위해서 간단한 파이썬 코드를 작성해서 진행하였고, 해당 코드는 아래와 같습니다.

import os
import random
import shutil

# 원본 이미지 파일이 들어있는 폴더 경로
source_folder = "원본 파일 경로"

# 셔플된 이미지 파일 리스트 생성
image_list = [f for f in os.listdir(source_folder) if f.endswith('.jpg')]
random.shuffle(image_list)

# 이동할 이미지 개수
num_to_move = 1000

# 이동할 이미지 파일 경로 리스트 생성
move_list = image_list[:num_to_move]

# 이미지 파일 이동할 폴더 경로
target_folder = "이동할 경로" # 새로 만든 폴더 train/val/test

# 이미지 파일 이동
for file_name in move_list:
    # 원본 이미지 파일 경로
    source_file_path = os.path.join(source_folder, file_name)
    # 대상 폴더 내 이동할 이미지 파일 경로
    target_file_path = os.path.join(target_folder, file_name)
    shutil.move(source_file_path, target_file_path)

print(f"{num_to_move} images moved.")

 추가적으로 해당 이미지가 있는 각각의 폴더에 맞는 라벨이 들어가야 하므로 아래와 같은 파이썬 코드를 통해 텍스트 파일 형태의 라벨도 옮겨주었습니다.

import os
import shutil

# 이미지 파일과 텍스트 파일이 저장된 디렉토리 경로
image_dir = "이미지 파일 경로" # 옮겨진 이미지가 있는 경로
text_dir = "텍스트 파일 경로" # 옮겨지지 않은 텍스트 파일이 있는 경로
dest_dir = "옮겨질 경로" # 텍스트 파일이 옮겨져야 하는 경로

# 이미지 파일 이름 목록을 작성합니다.
image_filenames = os.listdir(image_dir)

# 이미지 파일 이름에서 확장자를 분리합니다.
image_names = [os.path.splitext(filename)[0] for filename in image_filenames]

# 텍스트 파일 이름 목록을 작성합니다.
text_filenames = os.listdir(text_dir)

# 이미지 파일 이름과 일치하는 텍스트 파일을 찾아서 이동시킵니다.
for image_name in image_names:
    text_filename = os.path.join(text_dir, image_name + ".txt")
    if os.path.exists(text_filename):
        shutil.move(text_filename, os.path.join(dest_dir, os.path.basename(text_filename)))

 해당 과정이 완료되면 아래 사진과 같이 파일이 만들어지고 test와 val의 images 폴더에 각각 1,000장의 이미지와 labels 폴더에 1,000개의 파일이 옮겨지게 됩니다. 또한 train의 images에는 8,000장의 이미지와 labels에는 8,000개의 텍스트 파일이 존재합니다.

데이터셋 나누기

4. 모델 학습하기

 이전까지의 내용은 YOLOv5 모델을 통해 학습을 하기 위한 데이터셋을 준비하는 과정이었습니다. 파이참을 이용해서 모델 학습을 진행하기 위한 과정은 아래와 같습니다.

 1) 우선 파이참을 열고 Get from VCS를 선택합니다.

2) YOLOv5의 깃허브 주소를 URL에 넣고 Clone 버튼을 선택합니다.

YOLOv5의 깃허브 주소는 다음과 같습니다.

https://github.com/ultralytics/yolov5

GitHub - ultralytics/yolov5: YOLOv5 🚀 in PyTorch > ONNX > CoreML > TFLite

3) 앞선 과정에서 생성된 경로로 이동해서 해당 경로에 준비해둔 custom dataset 폴더를 복사합니다. 이후 yolov5 폴더 내 data 폴더에 custom_dataset.yaml 파일을 아래와 같이 생성합니다.

path: C:\Users\admin\PycharmProjects\yolov5\custom_dataset #custom_dataset 폴더 경로
train: train\images
val: val\images
test: test\images  

# Classes
nc : 2 # 분류할 class 갯수
names: ['normal', 'butane'] # 분류할 class 이름

4) 준비가 완료되면 하단 Terminal의 Command Prompt에 들어가서 하단 이미지와 같이

python train.py --data custom_dataset.yaml --epochs 100 을 입력한 뒤 엔터를 누르고 훈련이 완료되기를 기다립니다. 해당 프로젝트에서는 epochs을 100으로 설정하는 것 이외에는 모두 default 값으로 설정했습니다.

 

train 진행 중

 GPU가 없는 환경에서 프로젝트를 진행하였고, 총 100회의 훈련이 완료되는데 약 126시간이 소요되었습니다.

train 완료

5) 훈련과정이 완료되면 학습한 모델을 검증하는 과정을 거치게 됩니다. 검증을 위해서는 다음과 같은 내용을 입력합니다.

python val.py --data custom_dataset.yaml --weights .\runs\train\exp\weights\best.pt

이번 프로젝트에서 검증 과정에는 데이터 경로와 모델 가중치외에는 전부 default값을 사용하였습니다. 해당 과정은 짧은 시간 안에 마무리 됩니다.

val

6) 마지막으로 학습한 모델을 통해 예측이 정상적으로 수행되는지 확인합니다. 이때 앞서 분류했던 test 폴더에 담긴 이미지들이 사용됩니다.

python detect.py --weights .\runs\train\exp\weights\best.pt --source .\custom_dataset\test\images

test

 해당 내용은 미리 분류한 test 이미지 뿐만 아니라 웹캠이나 영상 파일에서도 사용할 수 있습니다.

 

7) 해당 과정이 전부 다 완료되면 새롭게 생성된 runs폴더의 detect-exp 경로로 들어가면 test의 images 폴더에 있던 이미지들에 대한 결과를 확인할 수 있게 됩니다.

 

예측 결과