본문 바로가기

IT/Deeplearning

[YOLO / Object Detection / Keras] Code Review - [1]

이번 포스팅에서는 Keras로 짜여진 YOLO Darknet 코드에 대해서 코드 리뷰를 진행하려고합니다.

Keras 코드를 뜯어보는 이유는 기존 YOLO 프로젝트가 C로 짜여져있어서, 직관적인 컨셉을 이해하기 많이 어렵기 때문에

전반적인 핵심 컨셉을 Keras YOLO의 코드 리뷰를 통해서 파악하고, YOLO를 기반으로 코드를 변경하는 프로젝트를 진행하면서

기존 YOLO 프로젝트를 변경하는데 도움을 주려고하는 이유입니다.




1. 코드


코드는 제 깃헙의 keras-yolo 프로젝트를 기반으로 합니다.

VOC Pascal 데이터로 직접 학습하고 해당 결과를 확인했으므로 실행하실 분들은 실행시켜보셔도 됩니다.




2. Train 코드


먼저 학습코드에 대해서 살펴보겠습니다.

YOLO Darknet을 학습시키기 위한 Config.json파일은 다음과 같습니다.


config.json

{
    "model" : {
        "architecture":         "Full Yolo",
        "input_size":           416,
        "anchors":              [0.57273, 0.677385, 1.87446, 2.06253, 3.33843, 5.47434, 7.88282, 3.52778, 9.77052, 9.16828],
        "max_box_per_image":    10,        
        "labels":               ["aeroplane", "bicycle", "bird", "boat", "bottle", "bus", "car", "cat", "chair", "cow", "diningtable", "dog", "horse", "motorbike", "person", "pottedplant", "sheep", "sofa", "train", "tvmonitor"]

    },

    "train": {
        "train_image_folder":   "/media/martin/enumcut/EnumNet/EnumCut/DataSet/VOC/VOCdevkit/VOC2012/JPEGImages/",
        "train_annot_folder":   "/media/martin/enumcut/EnumNet/EnumCut/DataSet/VOC/VOCdevkit/VOC2012/Annotations/",
          
        "train_times":          10,
        "pretrained_weights":   "None",
        "batch_size":           16,
        "learning_rate":        1e-4,
        "nb_epoch":             50,
        "warmup_epochs":        3,

        "object_scale":         5.0 ,
        "no_object_scale":      1.0,
        "coord_scale":          1.0,
        "class_scale":          1.0,

        "saved_weights_name":   "yolo_voc_2007.h5",
        "debug":                true
    },

    "valid": {
        "valid_image_folder":   "",
        "valid_annot_folder":   "",

        "valid_times":          1
    }
}


다음으로 train.py를 보도록 하겠습니다.


train.py

import argparse
import os
import numpy as np
from preprocessing import parse_annotation
from frontend import YOLO
import json

os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"]="0"

argparser = argparse.ArgumentParser(
    description='Train and validate YOLO_v2 model on any dataset')

argparser.add_argument(
    '-c',
    '--conf',
    help='path to configuration file')

def _main_(args):

    config_path = args.conf

    # JSON 파일 로드
    with open(config_path) as config_buffer:
        config = json.loads(config_buffer.read())

    ###############################
    #   Parse the annotations 
    ###############################

    # parse annotations of the training set
    train_imgs, train_labels = parse_annotation(config['train']['train_annot_folder'], 
                                                config['train']['train_image_folder'], 
                                                config['model']['labels'])
.
.
.
(생략)

if __name__ == '__main__':
    args = argparser.parse_args()
    _main_(args)

1. argparser에 대한 description과 confg 옵션을 설정한다.

2. 학습을 걸 때, 명령어는 python train.py -c config.json이므로 -c config.json에 대한 파라미터를 파싱한 Object를 _main_함수에 넘긴다.

3. config_path에 받은 config path인 config.json을 넘긴다.

4. config.json을 읽어들인다.

5. parse_annotation 메소드에 읽어들인 config.json파일의 학습 레이블 경로와 학습 이미지 경로, 그리고 클래스 리스트를 파라미터로 넘겨서 train_imgs, train_labels를 받습니다.


preprecessing.py

def parse_annotation(ann_dir, img_dir, labels=[]):
    all_imgs = []
    seen_labels = {}
    
    for ann in sorted(os.listdir(ann_dir)):
        img = {'object':[]}
        
        tree = ET.parse(ann_dir + ann)
        
        for elem in tree.iter():
            if 'filename' in elem.tag:
                img['filename'] = img_dir + elem.text
            if 'width' in elem.tag:
                img['width'] = int(elem.text)
            if 'height' in elem.tag:
                img['height'] = int(elem.text)
            if 'object' in elem.tag or 'part' in elem.tag:
                obj = {}
                
                for attr in list(elem):
                    if 'name' in attr.tag:
                        obj['name'] = attr.text

                        if obj['name'] in seen_labels:
                            seen_labels[obj['name']] += 1
                        else:
                            seen_labels[obj['name']] = 1
                        
                        if len(labels) > 0 and obj['name'] not in labels:
                            break
                        else:
                            img['object'] += [obj]
                            
                    if 'bndbox' in attr.tag:
                        for dim in list(attr):
                            if 'xmin' in dim.tag:
                                obj['xmin'] = int(round(float(dim.text)))
                            if 'ymin' in dim.tag:
                                obj['ymin'] = int(round(float(dim.text)))
                            if 'xmax' in dim.tag:
                                obj['xmax'] = int(round(float(dim.text)))
                            if 'ymax' in dim.tag:
                                obj['ymax'] = int(round(float(dim.text)))

        if len(img['object']) > 0:
            all_imgs += [img]
                        
    return all_imgs, seen_labels

1. all_imgs라는 빈 리스트를 만듭니다.

2. ssen_labels라는 딕셔너리를 만듭니다.

3. ann_dir의 경로 내부 파일들의 리스트를 읽어 들여와서 정렬을 하고 차례차례 순회를 돈다.

4. img이라는 딕셔너리에 "ojbect"라는 빈 리스트를 갖는 요소를 추가한다.

5. xml파일을 읽어들인다.

6. 탑다운방식으로 iteration을 돌면서 tag object를 순회한다.

6-1. tag안에 filename이라는 것이 있으면 딕셔너리 img에 filename와 이미지 경로가 쌍을 이루도록 저장한다.

6-2. tag안에 width라는 것이 있으면 딕셔너리 img에 width와 width 값이 쌍을 이루도록 저장한다.

6-3. tag안에 height라는 것이 있으면 딕셔너리 img에 height와 height 값이 쌍을 이루도록 저장한다.

6-4. tag안에 part나 object라는 것이 있는지 찾는다.

6-5. obj라는 빈 딕셔너리를 생성한다.

6-6. 리스트 elem의 속성을 가지고 순회를 돈다

6-5-1. tag 속성에 name이 있으면 딕셔너리 obj에 name과 name의 값이 쌍을 이루도록 저장한다.

6-5-2. 만약에 obj의 name속성의 값이 seen_labels안에 있으면 해당 카운트를 1 올린다.

6-5-3. 만약에 obj의 name속성의 값이 seen_labels안에 없으면, seen_labels안에 obj의 name 속성값을 추가하고 값을 1로 초기화한다.(label counting)

6-5-4. 만약 labels의 길이가 0보다 크면서 obj의 속성 name이 labels안에 없다면 순회룰 중지한다.

6-5-5. 그렇지 않다면 딕셔너리 img의 object 속성안에 obj값을 추가한다.

6-5-6. tag 속상값에 bndbox가 있으면, 딕셔너리 obj안에 xmin, ymin, xmax, ymax 값을 xml값에서 가져와서 대입해준다.

7.딕셔너리 img의 object 속성값이 0보다 크다면, all_imgs에 딕셔너리 img를 추가한다.

8. all_imgs, seen_labels를 반환한다.


결국 parse_annotation이라는 메소드는 object의 개수 총합과 각 object 종류마다의 갯수 그리고 각 object의 속성값들을 json형태와 비슷한 딕셔너리에 저장해주는 역활을 합니다.