google Open Image dataset 사용해 YOLO training / OID_v3_toolkit 사용하기

github: github.com/chaeyeongyoon/ComputerVision_Study

chaeyeongyoon/ComputerVision_Study

 

Google Image dataset은 bounding box가 있는 600여개의 object category 의 bounding box정보를 가진 이미지가 있습니다.

최근에는 training/validation/test set 뿐만 아니라 segmentation도 제공합니다.

한 이미지에 여러 object를 포함한 아주 복잡한 image data set입니다.

storage.googleapis.com/openimages/web/index.html

Open Images V6

현재 v6까지 release 되어있습니다. 홈페이지에 들어가면 다음과 같이 카테고리 정보와 dataset을 다운할 수 있습니다.

OIDv4_toolkit을 이용해 open Images v4의 원하는 카테고리의 이미지를 다운받았습니다.

OIDv4_toolkit

github.com/EscVM/OIDv4_ToolKit

EscVM/OIDv4_ToolKit

위의 github을 clone한 후 requirements를 설치해줍니다

(anaconda prompt에서 다음 명령어 실행)

# requirements 설치
pip install -r requirements.txt

# Football, Tennis Ball, Person 카테고리의 이미지들 (카테고리당 300개로 limit)을 공통폴더로 다운로드
# csv형태의 annotation 파일 다운
python main.py downloader --classes Football "Tennis Ball" Person -type_csv train --multiclasses 1 --limit 300

OID_v4_toolkit으로 다운받은 directory 구조

google open image의 csv파일에는 Image ID, Source, LabelName, Xmin, Xmax, Ymin, Ymax, IscOccluded, IsTruncated, IsGroupOf, IsDepiction, IsInside등의 정보가 포함되어있습니다. 그리고 다음의 option들은 다운받을때 설정 가능

]

각 이미지의 이름으로 된 txt파일들은 bounding box좌표정보를 소수점으로 가지고 있습니다.

따라서 이 파일들을 VOC xml파일 형태로 변경한 후 keras yolo용 csv파일로 변환시켜줄 것 입니다.

VOC xml파일 타입으로 변경하는 것은 oid_to_pascal_vod_xml.py를 이용합니다.

oid_to_pascal_vod_xml.py를 OID_v4_Toolkit 폴더에 옮겨준 후 실행시켜주면 OID폴더 밑의 모든 이미지들에 대해 변환해줍니다.

xml파일 예

 

이후 xml파일들을 다음 함수의 실행으로 (raccon dataset training과 유사합니다) csv파일로 변경시켜줍니다.

import glob
import xml.etree.ElementTree as ET

classes_map = {'Football':0, 'Tennis_ball':1, 'Person':2}

# voc xml -> yolo csv
def xml_to_csv(path, output_filename):
    xml_list = []
    # xml 확장자 가진 모든 파일의 절대경로로 xml_file할당
    with open(output_filename, 'w') as train_csv_file:
        for xml_file in glob.glob(path + '/*.xml'):
            # xml file parsing and make Element Tree - extract object information
            tree = ET.parse(xml_file)
            root = tree.getroot()
            full_image_name = os.path.join(IMAGE_DIR, root.find('filename').text)
            value_str_list = ' '
            for obj in root.findall('object'):
                xmlbox = obj.find('bndbox')
                class_name = obj.find('name').text
                x1 = int(xmlbox.find('xmin').text)
                y1 = int(xmlbox.find('ymin').text)
                x2 = int(xmlbox.find('xmax').text)
                y2 = int(xmlbox.find('ymax').text)
                class_id = classes_map[class_name]
                value_str = f'{x1},{y1},{x2},{y2},{class_id}'
                value_str_list += value_str + ' '
            train_csv_file.write(full_image_name + value_str_list + '\n')

이후 앞선 게시물에서와 같이 keras-yolo3 package다운 - 시스템경로추가 - pretrained weights파일 h5파일로 변환

수행해 준 후 training을 시작합니다.

 

<여러 경로들과 class, anchor정보를 정의>

from yolo3.model import preprocess_true_boxes, yolo_body, tiny_yolo_body, yolo_loss
from yolo3.utils import get_random_data

# class정보담은 txt파일 생성
BASE_DIR = os.path.join(HOME_DIR, 'DLCV/Detection/yolo/keras-yolo3')
classes_path = os.path.join(BASE_DIR, 'model_data/ballnperson_classes.txt')
with open(classes_path, 'w') as f:
    f.write('Football\nTennis_ball\nPerson')
    
from train import get_classes, get_anchors
from train import create_model, data_generator, data_generator_wrapper

BASE_DIR = os.path.join(HOME_DIR, 'DLCV/Detection/yolo/keras-yolo3')

## annotation file path, epoch model save path, object class file path, anchor file path
annotation_path = os.path.join(ANNO_DIR, 'ballnperson_anno.csv')
log_dir = os.path.join(BASE_DIR, 'snapshots/ballnperson/')
classes_path = os.path.join(BASE_DIR, 'model_data/ballnperson_classes.txt')
anchors_path = os.path.join(BASE_DIR, 'model_data/yolo_anchors.txt')

class_names = get_classes(classes_path)
num_classes = len(class_names)
anchors = get_anchors(anchors_path)
print(class_names, num_classes)
print(anchors)

<config class로 training에 필요한 변수들 깔끔하게 정의>

 

# read csv annotation file and convert to list ( lines )
with open(annotation_path) as f:
    lines = f.readlines()
    
class config:
    #initial weights
    initial_weights_path=os.path.join(BASE_DIR, 'model_data/yolo.h5' )
    # input_shape는 고정. 
    input_shape=(416, 416)
    # epochs는 freeze, unfreeze 2 step에 따라 설정. 
    first_epochs=50
    first_initial_epochs=0
    second_epochs=100
    second_initial_epochs=50
    # batch size, train,valid count, epoch steps   
    batch_size = 4
    val_split = 0.1   
    num_val = int(len(lines)*val_split)
    num_train = len(lines) - num_val
    train_epoch_steps = num_train//batch_size 
    val_epoch_steps =  num_val//batch_size
    
    anchors = get_anchors(anchors_path)
    class_names = get_classes(classes_path)
    num_classes = len(class_names)
    # epoch시 저장된 weight 파일 디렉토리 
    log_dir = os.path.join(BASE_DIR, 'snapshots/ballnperson/')
    
print('Class name:', config.class_names,'\nNum classes:', config.num_classes)

<csv 파일을 data_generator_wrapper에 넣어주어 train/valid generator만들어주는 함수>

<yolo model 형성 함수>

<callback함수 정의 함수>

정의

def create_generator(lines):
    
    train_data_generator = data_generator_wrapper(lines[:config.num_train], config.batch_size, 
                                                  config.input_shape, config.anchors, config.num_classes)
    
    valid_data_generator = data_generator_wrapper(lines[config.num_train:], config.batch_size, 
                                                  config.input_shape, config.anchors, config.num_classes)
    
    return train_data_generator, valid_data_generator
    
def create_yolo_model():
    is_tiny_version = len(config.anchors)==6 
    if is_tiny_version:
        model = create_tiny_model(config.input_shape, config.anchors, config.num_classes, 
            freeze_body=2, weights_path=config.initial_weights_path)
    else:
        model = create_model(config.input_shape, config.anchors, config.num_classes, 
            freeze_body=2, weights_path=config.initial_weights_path)
        
    return model
    
def create_callbacks():
    logging = TensorBoard(log_dir=config.log_dir)
    checkpoint = ModelCheckpoint(config.log_dir + 'ep{epoch:03d}-loss{loss:.3f}-val_loss{val_loss:.3f}.h5',
                                 monitor='val_loss', save_weights_only=True, save_best_only=True, period=3)
    reduce_lr = ReduceLROnPlateau(monitor='val_loss', factor=0.1, patience=3, verbose=1)
    early_stopping = EarlyStopping(monitor='val_loss', min_delta=0, patience=10, verbose=1)
    
    # return as list
    return [logging, checkpoint, reduce_lr, early_stopping]

 

<train>

train_data_generator, valid_data_generator = create_generator(lines)
ballnperson_model = create_yolo_model()
callback_list = create_callbacks()

print('### First train start ###')
ballnperson_model.compile(optimizer=Adam(lr=1e-3), loss={'yolo_loss':lambda y_true, y_pred: y_pred})
ballnperson_model.fit_generator(train_data_generator, steps_per_epoch=config.train_epoch_steps,
                                validation_data= valid_data_generator, validation_steps = config.val_epoch_steps,
                                epochs=config.first_epochs, initial_epoch=config.first_initial_epochs, callbacks=callback_list)

ballnperson_model.save_weights(log_dir + 'trained_weights_stage_1.h5')

# second train
for i in range(len(ballnperson_model.layers)):
    ballnperson_model.layers[i].trainable = True
    
print('### Second train start ###')
ballnperson_model.compile(optimizer=Adam(lr=1e-4), loss={'yolo_loss':lambda y_true, y_pred: y_pred})
ballnperson_model.fit_generator(train_data_generator, steps_per_epoch=config.train_epoch_steps,
                                validation_data= valid_data_generator, validation_steps = config.val_epoch_steps,
                                epochs=config.second_epochs, initial_epoch=config.second_initial_epochs, callbacks=callback_list)

ballnperson_model.save_weights(log_dir + 'trained_weights_final.h5')

이렇게 저장된 weights로 inference 진행해줍니다. (inference게시물 참조)