Table of Contents
  1. What is the YOLO?
    1. YOLO Architecture
    2. YOLO Output
  2. 학습 데이터 준비
    1. Image Gathering
    2. Image Labeling
    3. Image Resizing
    4. Image Generating
  3. YOLO Training
    1. Convert Darknet Model To Keras Model
    2. Convert Annotation
    3. Transfer Learning
  4. 마치며

실시간 객체 탐지 모델인 YOLO 모델을 학습하는 방법에 대해 공유합니다.

1부. 딥러닝으로 더욱 현실감 있는 AR 앱 만들기
2부. ARKit을 이용하여 iOS AR앱 만들기
✔︎ 3부. Custom YOLO v3 모델 만들기
4부. CoreML을 이용하여 iOS에 YOLO v3 모델 탑재하기

What is the YOLO?

시작하기에 앞서 YOLO에 대한 간단한 설명과 어떻게 Training 방법에 대해 알아보도록 하겠습니다.

✔︎ 진행 순서

  • YOLO Architecture & Output
  • Object Detection Data 준비
  • Transfer Learning

본 포스팅은 YOLO 모델 활용에 대한 글이므로, YOLO에 대한 자세한 설명은 다른 글을 참고하시기 바랍니다.

YOLO Architecture

YOLO (You Only Look Once)는 학습한 물체의 종류와 위치를 실시간으로 파악 할 수 있는 Real-Time Object Detection 모델로 다음과 같은 아키텍쳐로 구성되어 있습니다.

YOLO v3 Architecture<br>출처 - <a href=https://www.cyberailab.com/home/a-closer-look-at-yolov3>https://www.cyberailab.com</a>YOLO v3 Architecture
출처 - https://www.cyberailab.com

YOLO Output

YOLO 모델의 output은 N x N x (C + 5) x B) 형태의 텐서(tensor)를 반환하며, 이 정보를 이용하여 이미지 상에 존재하는 물체의 종류와 2D 좌표를 획득할 수 있습니다.

N x N x (C + 5) x B)

  • N x N : Grid
  • C : 학습한 Object 수(Classes)
  • 5 : 탐지한 Object의 Bounding Box 정보(x, y, w, h, confidence)
  • B : Grid cell당 박스 수

YOLO output<br>출처 - <a href=https://www.cyberailab.com/home/a-closer-look-at-yolov3>https://www.cyberailab.com</a>YOLO output
출처 - https://www.cyberailab.com

다음으로 YOLO 모델을 학습을 위한 데이터 준비 방법을 알아보도록 하겠습니다.

학습 데이터 준비

Object Detection 모델을 만들기 위해서는 오브젝트가 담긴 이미지와 이미지 상의 오브젝트의 위치와 종류를 나타내는 라벨이 필요합니다. 이번 프로젝트에서는 다음과 같은 순서로 데이터를 준비하였습니다.

✔︎ 데이터 준비 순서

  • Image Gathering
  • Image Labeling
  • Image Resizing
  • Image Generating

Image Gathering

이번 프로젝트에서는 램프, 전구, 형광등 총 3개의 물체를 학습하였고 이를 위해 다음과 같은 이미지를 수집하였습니다. 램프와 전구의 경우 ON/OFF 상태를 구분할 수 있도록 다른 클래스로 분리하였기 때문에 총 클래스 수는 5개입니다.

조명 인식 모델을 만들기 위해 사용된 이미지조명 인식 모델을 만들기 위해 사용된 이미지

동영상으로 해당 오브젝트들을 촬영한 뒤 각 프레임들을 이미지로 변환하는 방법으로 총 5천장의 이미지를 수집하였습니다.

OpenCV를 이용한 video to image
1
2
3
4
5
6
7
8
9
10
11
12
import cv2
for i in range(len(videoFiles)):
    cam = cv2.VideoCapture(videoFile)
    currentFrame = 0
    while(True):
        ret, frame = cam.read()
        if ret:
            cv2.imwrite(currentFrame + '.jpg', frame)
            currentFrame += 1
        else:
            break
    cam.release()

Image Labeling

Object Detection 모델을 만들기 위해서는 학습하고자 하는 물체가 담긴 이미지와 해당 물체의 종류 그리고 위치 정보가 기록된 label이 필요합니다. label의 경우 직접 수작업으로 데이터를 생성해야 하는데, 작업을 편하게 하기 위해 labelImg라는 오픈소스 툴을 이용하였습니다.

labelImg는 이미지 상의 오브젝트의 위치와 종류를 xml 형태로 반환합니다.
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
<annotation>
	<folder>lamp_on</folder>
	<filename>0.jpg</filename>
	<path>/lamp_on/0.jpg</path>
	<source>
		<database>Unknown</database>
	</source>
	<size>
		<width>1920</width>
		<height>1080</height>
		<depth>3</depth>
	</size>
	<segmented>0</segmented>
	<object>
		<name>lamp_on</name>
		<pose>Unspecified</pose>
		<truncated>0</truncated>
		<difficult>0</difficult>
		<bndbox>
			<xmin>704</xmin>
			<ymin>384</ymin>
			<xmax>970</xmax>
			<ymax>546</ymax>
		</bndbox>
	</object>
</annotation>

Image Resizing

수집한 이미지의 원본은 1920 x 1080로 학습 시간이 많이 소요되기 때문에, Pillow 라이브러리를 이용하여 이미지의 크기를 10분의 1 크기인 192 x 108로 줄였습니다.

image resizing
1
2
3
4
5
from PIL import Image
for image_file in images:
  image = Image.open(image_file)
  resize_image = image.resize((192, 108))
  resize_image.save(new_path)
label resizing
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
def changeLabel(xmlPath, newXmlPath, imgPath, boxes):
    tree = elemTree.parse(xmlPath)

    # path 변경
    path = tree.find('./path')
    path.text = imgPath[0]

    # bounding box 변경
    objects = tree.findall('./object')
    for i, object_ in enumerate(objects):
        bndbox = object_.find('./bndbox')
        bndbox.find('./xmin').text = str(boxes[i][0])
        bndbox.find('./ymin').text = str(boxes[i][1])
        bndbox.find('./xmax').text = str(boxes[i][2])
        bndbox.find('./ymax').text = str(boxes[i][3])
    tree.write(newXmlPath, encoding='utf8')

Image Generating

양질의 데이터가 많을수록 모델의 성능은 좋아지지만, 너무 많은 데이터는 직접 라벨링을 하기가 힘들기 때문에 라벨링을 완료한 5천장의 이미지를 이용하여 추가 이미지를 생성하였습니다. 이미지를 생성하는 다양한 방법들이 있지만, rotation의 경우 Bounding Box의 표기법의 한계로 인해 오차가 커지기 때문에, horizontal flip을 이용한 방법으로 이미지를 1만장으로 늘렸습니다.

horizontal flip
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
import random
import numpy as np
class RandomHorizontalFlip(object):
    def __init__(self, p=0.5):
        self.p = p

    def __call__(self, img, bboxes):
        img_center = np.array(img.shape[:2])[::-1]/2
        img_center = img_center.astype(int)
        img_center = np.hstack((img_center, img_center))
        if random.random() < self.p:
            img = img[:, ::-1, :]
            bboxes[:, [0, 2]] += 2*(img_center[[0, 2]] - bboxes[:, [0, 2]])
            box_w = abs(bboxes[:, 0] - bboxes[:, 2])
            bboxes[:, 0] -= box_w
            bboxes[:, 2] += box_w
        return img, bboxes
image와 함께 label도 generating 하고있습니다.
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
for imgFile in imgFiles:
    fileName = imgFile.split('.')[0]
    label = f'{labelPath}{fileName}.xml'
    w, h = getSizeFromXML(label)

    # opencv loads images in bgr. the [:,:,::-1] does bgr -> rgb
    image = cv2.imread(imgPath + imgFile)[:,:,::-1]
    bboxes = getRectFromXML(classes, label)

    # HorizontalFlip image
    image, bboxes = RandomHorizontalFlip(1)(image.copy(), bboxes.copy())

    # Save image
    image = Image.fromarray(image, 'RGB')
    newImgPath = f'./data/light/image/train/{className}/'
    if not os.path.exists(newImgPath):
        os.makedirs(newImgPath)
    image.save(newImgPath + imgFile)

    # Save label
    newXmlPath = f'./data/light/label/train/{className}/'
    if not os.path.exists(newXmlPath):
        os.makedirs(newXmlPath)
    newXmlPath = newXmlPath + fileName + '.xml'
    changeLabel(label, newXmlPath, newImgPath, bboxes)

이렇게 라벨링된 이미지 준비를 끝마쳤습니다. 다음으로 데이터를 이용하여 YOLO 모델을 학습시키는 방법에 대해 설명하도록 하겠습니다.

YOLO Training

Keras-yolo3 라이브러리를 이용하여 YOLO 모델을 학습시켰습니다.

✔︎ 진행순서

  • Convert Darknet Model To Keras Model
  • Convert Annotation
  • Training

Convert Darknet Model To Keras Model

YOLO v3 모델의 경우 모바일 기기에 올리기에는 너무 무거움으로 경량화 버전인 Tiny YOLO v3 모델을 이용하였고, 기존에 학습된 weight를 그대로 사용할 수 있는 전이학습(transfer learning)을 이용하여 학습을 진행하였습니다.

tiny yolov3 pretrained weights 다운로드
1
wget https://pjreddie.com/media/files/yolov3-tiny.weights

위 명령어를 실행하면 yolov3-tiny.weights라는 35.4MB의 weights를 다운받을 수 있습니다. YOLO3 자체는 C/C++로 구현된 DarkNet 프레임워크로 구현되어 있습니다. 이를 Keras에서 사용할 수 있는 .h5 포멧으로 변환해야합니다. 변환 방법은 Keras-yolo3의 convert.py 파일을 이용합니다.

Convert darknet model to keras model
1
python convert.py yolov3-tiny.cfg yolov3-tiny.weights model_data/yolo_tiny.h5

✔︎ Convert darknet model to keras model

  • convert.py : 변환명령 수행 파일
  • yolov3.cfg : Darknet에서 사용하는 모델 구조 정의 파일
  • yolov3.weight : Darknet으로 학습된 모델 파일

완료가 되면 .h5형태로 weights 파일이 변환됩니다. 다음으로 변환된 모델이 잘 작동하는지 테스트를 위해 아래 샘플 이미지를 이용해보도록 하겠습니다.

sample image : dog.jpg<br>출처 : <a href=https://github.com/pjreddie/darknet>darknet</a>sample image : dog.jpg
출처 : darknet
tiny YOLO v3 converted model test
1
2
3
4
5
6
7
8
9
10
11
from IPython.display import display
from PIL import Image
from yolo import YOLO

def objectDetection(file, model_path, class_path):
    yolo = YOLO(model_path=model_path, classes_path=class_path, anchors_path="model_data/tiny_yolo_anchors.txt")
    image = Image.open(file)
    result_image = yolo.detect_image(image)
    display(result_image)

objectDetection('dog.jpg', 'model_data/yolo_tiny.h5', 'model_data/coco_classes.txt')

YOLO 모델이 적절한 bounding box와 confidence 그리고 class를 반환하는 것을 보니 변환에 성공한 것 같습니다.

자전거, 개, 자동차를 탐지한 YOLO model<br>출처 : <a href=https://github.com/pjreddie/darknet>darknet</a>자전거, 개, 자동차를 탐지한 YOLO model
출처 : darknet

Convert Annotation

Keras-yolo3에서 데이터를 학습하기 위해선 위에서 labeling한 데이터를 아래와 같은 형태로 변환해야합니다. voc_annotation.py 파일을 참고하여 아래와 같은 형태로 변환하시면 됩니다.

✔︎ Annotation Format

  • Row format: image_file_path box1 box2 … boxN
  • Box format: x_min,y_min,x_max,y_max,class_id (no space)
Annotation example
1
2
path/to/img1.jpg 50,100,150,200,0 30,50,200,120,3
path/to/img2.jpg 120,300,250,600,2
Convert annotation
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
import xml.etree.ElementTree as ET
from os import getcwd
import glob

def convert_annotation(annotation_voc, train_all_file):
    tree = ET.parse(annotation_voc)
    root = tree.getroot()

    for obj in root.iter('object'):
        difficult = obj.find('difficult').text
        cls = obj.find('name').text
        if cls not in classes or int(difficult)==1: continue
        cls_id = classes.index(cls)
        xmlbox = obj.find('bndbox')
        b = (int(xmlbox.find('xmin').text), int(xmlbox.find('ymin').text), int(xmlbox.find('xmax').text), int(xmlbox.find('ymax').text))
        train_all_file.write(" " + ",".join([str(a) for a in b]) + ',' + str(cls_id))

train_all_file = open('./data/light/train_all.txt', 'w')

# Get annotations_voc list
for className in classes:
    annotations_voc = glob.glob(f'./data/light/label/train/{className}/*.xml')
    for annotation_voc in annotations_voc:
        image_id = annotation_voc.split('/')[-1].split('.')[0]+'.JPG'
        train_all_file.write(f'./data/light/image/train/{className}/{image_id}')
        convert_annotation(annotation_voc, train_all_file)
        train_all_file.write('\n')
train_all_file.close()
필자의 annotation 변환 결과
1
2
3
4
./data/light/image/train/bulb_on/162.JPG 70,0,124,52,0
./data/light/image/train/bulb_on/1390.JPG 70,61,117,100,0
./data/light/image/train/bulb_on/604.JPG 78,38,93,52,0
./data/light/image/train/bulb_on/88.JPG 72,28,139,93,0

Transfer Learning

마지막으로 학습은 train.py 파일의 _main() 함수에 다음과 같은 변수가 있는데 이를 수정하여 실행하면 됩니다.

train.py의 _main()
1
2
3
4
5
6
7
def _main():
  annotation_path = 'data/light/train_all.txt' # 위 Convert Annotation의 결과 파일경로
  log_dir = 'logs/000/' # 학습 로그 저장경로
  classes_path = 'data/light/classes.txt' # 학습하는 클래스 목록
  anchors_path = 'model_data/yolo_tiny_anchors.txt' # tiny YOLO 모델의 anchors
  transfer_learning_path = 'model_data/yolo_tiny.h5' # Convert Darknet Model To Keras Model 결과 h5파일
  # ... 이하 생략

위와같이 수정을 한 뒤 실행을 하면 기존 pretraining 된 tiny YOLO weights를 로드한 뒤, 전이학습(Transfer Learning)을 시작합니다. 총 100epoch을 batch size 32로 진행하는데, 50에폭까지는 전체 44개의 Layer 중 42개를 Freeze 한 뒤 학습이 진행되고, 이후에는 모든 레이어를 Unfreeze한 뒤 학습이 진행됩니다.

하지만 Keras의 EarlyStop이 다음과 같은 조건으로 적용되어있어 전체 100epoch을 진행하기 전에 학습이 끝나기도 합니다.

EarlyStopping
1
EarlyStopping(monitor='val_loss', min_delta=0, patience=10, verbose=1)

또한 최적의 모델을 선택하기 위해 ModelCheckpoint 함수가 validation error 를 모니터링하면서, 이전 epoch 에 비해 validation performance 가 좋은 경우, 무조건 이 때의 parameter들을 중간중간 저장되고 있어서 학습이 도중에 멈추더라도 중간 결과물을 보존할 수 있습니다.

모든 학습이 완료되었으니 모델을 테스트하도록 하겠습니다.

Trained Model Test
1
2
3
4
5
6
7
8
9
10
11
from IPython.display import display
from PIL import Image
from yolo import YOLO

def objectDetection(file, model_path, class_path):
    yolo = YOLO(model_path=model_path, classes_path=class_path, anchors_path='model_data/yolo_tiny_anchors.txt')
    image = Image.open(file)
    result_image = yolo.detect_image(image)
    display(result_image)

objectDetection('data/light/test/562.jpg', 'model_data/light_tiny_model.h5', 'data/light/classes.txt')
전이학습이 완료된 tiny YOLO v3 output전이학습이 완료된 tiny YOLO v3 output

이미지에 담긴 조명의 정확한 Bounding Box와 신뢰도 그리고 종류를 잘 탐지하고 있습니다. 테스트까지 완료된 모델을 파일로 저장하기 위해 Keras-yolo3의 yolo.py 파일에 학습된 모델을 반환하는 함수를 추가하였습니다.

yolo.py 파일에 학습된 모델을 반환하는 함수를 추가
1
2
def get_model(self):
    return self.yolo_model

그리고 모델의 save함수를 이용하여 학습이 완료된 모델을 저장하였습니다.

save keras h5 model
1
2
model = yolo.get_model()
model.save('model_data/light_tiny_model.h5')

마치며

다음 4부에서는 학습한 YOLO Model을 iOS 어플리케이션에 올리는 방법에 대해 설명하도록 하겠습니다.

추가로, 많은 분들이 소스코드를 요청하여 github link를 공유드립니다.

1부. 딥러닝으로 더욱 현실감 있는 AR 앱 만들기
2부. ARKit을 이용하여 iOS AR앱 만들기
✔︎ 3부. Custom YOLO v3 모델 만들기
4부. CoreML을 이용하여 iOS에 YOLO v3 모델 탑재하기