YOLO 모델의 전이 학습(Transfer Learning), 클래스 추가 및 제거를 통한 커스텀 객체 검출 모델 구축

딥러닝 기술의 발달과 함께 객체 검출 기술도 비약적으로 발전하였습니다. 그 중에서 YOLO(You Only Look Once) 모델은 실시간 객체 검출을 가능하게 해주며, 다양한 분야에서 활용되고 있습니다. 특히 전이 학습(Transfer Learning)을 통해 개발자는 기존 모델을 기반으로 새로운 데이터를 학습시켜 쉽게 커스텀 객체 검출 모델을 구축할 수 있습니다. 이 글에서는 YOLO 모델의 전이 학습, 클래스 추가 및 제거를 통한 커스텀 객체 검출 모델 구축 방법에 대해 자세히 설명하겠습니다.

1. YOLO 모델 소개

YOLO(You Only Look Once)는 CNN(Convolutional Neural Network) 기반의 객체 검출 알고리즘으로, 이미지를 격자(grid)로 나누고 각 격자에서 객체를 동시에 예측하는 방식입니다. YOLO의 가장 큰 장점은 실시간으로 객체 검출을 수행할 수 있다는 것이며, 높은 정확도와 속도를 어느 정도로 함께 제공하는 것이 특징입니다.

2. 전이 학습(Transfer Learning)

전이 학습은 이미 학습된 모델을 새로운 작업에 재사용하는 기법입니다. 주로 대량의 데이터가 부족할 때, 사전 학습된 모델을 사용하여 적은 데이터로 더 나은 성능을 얻을 수 있게 해줍니다. YOLO 모델을 활용한 전이 학습은 여러 종류의 데이터에 적용하기 수월하며, 개별 객체에 대한 학습을 효율적으로 처리할 수 있습니다. 이번 섹션에서는 YOLO v3 모델을 사용하여 전이 학습을 수행하는 방법을 안내하겠습니다.

2.1 YOLO v3 활용

YOLO v3는 다양한 크기의 객체를 높은 정확도로 감지할 수 있는 능력을 갖추고 있습니다. 모델을 위해 Darknet이라는 오픈소스 프레임워크를 사용합니다. YOLO v3 모델을 설치하기 위해서는 다음과 같은 단계를 따라야 합니다.

2.1.1 환경 설정

우리는 darknet 프레임워크를 설치하여 YOLO v3를 사용할 것입니다. 아래의 명령어로 설치합니다:

git clone https://github.com/AlexeyAB/darknet.git 
cd darknet 
make

2.1.2 사전 학습된 모델 다운로드

YOLO v3에서 사전 학습된 가중치를 다운로드합니다. 이를 통해 전이 학습을 쉽게 수행할 수 있습니다.

wget https://pjreddie.com/media/files/yolov3.weights

3. 커스텀 데이터셋 구축

커스텀 객체 검출 모델을 구축하기 위해서는 우선 검출하고자 하는 객체의 이미지와 그에 대한 라벨이 필요합니다. 다음 단계로 커스텀 데이터셋을 준비하는 방법을 소개하겠습니다.

3.1 데이터 수집 및 라벨링

확인하고자 하는 특정 클래스의 이미지를 수집합니다. 예를 들어, ‘고양이’와 ‘강아지’를 검출하고 싶다면, 해당 이미지를 수집한 후 라벨링을 수행해야 합니다. 일반적으로 사용되는 라벨링 툴로는 LabelImg가 있습니다. 아래의 명령어로 설치할 수 있습니다.

git clone https://github.com/tzutalin/labelImg.git 
cd labelImg 
pip install -r requirements/requirements-linux-python3.txt 
make qt5py3

3.2 데이터셋 구조화

라벨링이 완료된 후, YOLO가 요구하는 형태로 데이터셋을 구성해야 합니다. 구조는 다음과 같습니다:

dataset/
    ├── images/
    │   ├── image1.jpg
    │   ├── image2.jpg
    │   └── ...
    └── labels/
        ├── image1.txt
        ├── image2.txt
        └── ...

4. YOLO 구성 파일 수정

YOLO 모델의 학습을 위해서는 몇 가지 구성 파일을 수정해야 합니다:

4.1 .cfg 파일 수정

YOLO v3의 .cfg 파일을 열고 클래스 수와 필터를 수정합니다. 예를 들어, 우리가 ‘고양이’와 ‘강아지’ 두 클래스를 검출할 경우, yolo 레이어의 필터 수는 다음과 같이 설정해야 합니다:

[yolo]
    ...
    filters=18  ; (classes + 5) * 3 (고양이 + 강아지 2) 
    classes=2

4.2 .data 파일 생성

다음으로, YOLO 모델이 사용할 데이터셋에 대한 정보를 담고 있는 .data 파일을 생성합니다:

classes= 2
    train  = data/train.txt
    valid  = data/valid.txt
    names = data/obj.names
    backup = backup/

4.3 객체 이름 파일 생성

검출하고자 하는 객체의 이름을 담고 있는 obj.names 파일을 생성합니다.

고양이
    강아지

5. 모델 학습

모든 준비가 완료되었으므로 이제 모델을 학습할 차례입니다. YOLO는 명령어를 통해 쉽게 학습할 수 있습니다. 다음의 명령어를 사용합니다:

./darknet detector train data/obj.data cfg/yolov3_custom.cfg yolov3.weights

6. 검증 및 테스트

모델 학습 후, 최종적인 성능 평가를 위해 검증 세트를 사용합니다. 학습이 끝난 모델은 다음 명령어로 테스트할 수 있습니다:

./darknet detector test data/obj.data cfg/yolov3_custom.cfg backup/yolov3_custom_final.weights data/test.jpg

7. 결과 분석

최종적으로 얻어진 검출 결과를 기준으로 정확도 및 성능을 평가합니다. 다양한 성능 지표를 사용하여 모델의 성능을 비교하고 분석합니다. 또한 다양한 시나리오를 시뮬레이션하여 향후 개선해야 할 사항을 도출합니다.

8. 결론

YOLO 모델은 객체 검출 분야에서 강력한 성능을 발휘하고 있습니다. 특히 전이 학습을 활용하여 적은 데이터로도 효과적인 학습을 진행할 수 있고, 클래스의 추가 및 제거를 통해 다양한 응용분야에 맞게 커스텀 모델을 구축할 수 있습니다. 이 글을 통해 독자들이 YOLO 모델을 활용하여 자신만의 객체 검출 모델을 손쉽게 만들 수 있기를 바랍니다.

참고자료

YOLO 및 전이 학습 관련 자료: