IBM_AI/vi-opton/opton_cls_train/make_data_files.py

import os
import argparse
import random
from config import config

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--data_dir', help='dataset dir',
                        default='./data', type=str)
    parser.add_argument('--train_percent', help='percent for train images (train+val = all * train_percent)',
                        default=0.9, type=float)
    parser.add_argument('--val_percent',  help='percent for val images (val = (train+val) * val_percent)',
                        default=0.1, type=float)
    args = parser.parse_args()
    return args

def remove_lastline_blank(path):
    f = open(path, 'rb+')
    if len(f.readlines()) > 0:
        f.seek(-1, os.SEEK_END)
        if f.next() == '\n':
            f.seek(-1, os.SEEK_END)
            f.truncate()
    f.close()

if __name__ == '__main__':
    args = parse_args()

    assert os.path.exists(args.data_dir), args.data_dir + 'does not exist'

    txt_dir = args.data_dir
    print('Save splits files to ' + os.path.abspath(txt_dir))

    cfg = config()
    class_idxes = cfg.get_model_index()

    # output files
    info_path = os.path.join(txt_dir, 'info.txt')
    train_path = os.path.join(txt_dir, 'train.txt')
    val_path = os.path.join(txt_dir, 'val.txt')
    train_val_path = os.path.join(txt_dir, 'trainval.txt')
    test_path = os.path.join(txt_dir, 'test.txt')

    info_file = open(info_path, 'w')
    train_file = open(train_path, 'w')
    val_file = open(val_path, 'w')
    train_val_file = open(train_val_path, 'w')
    test_file = open(test_path, 'w')

    total_all = 0
    total_train_val = 0
    total_val = 0
    total_train = 0
    total_test = 0

    # load all images
    dataset_dir = os.path.join(args.data_dir, 'dataset')
    class_dirs = os.listdir(dataset_dir)
    for class_dir in class_dirs:
        if class_dir not in class_idxes:
            print('[WARN] {} not in train classes'.format(class_dir))
            continue

        files_all = os.listdir(os.path.join(dataset_dir, class_dir))
        files_all = [n for n in files_all if n.endswith('.jpg')]

        num_all = len(files_all)
        num_train_val = int(num_all * args.train_percent)
        num_val = int(num_train_val * args.val_percent)
        num_train = num_train_val - num_val
        num_test = num_all - num_train_val
        files_train_val = random.sample(files_all, num_train_val)
        files_val = random.sample(files_train_val, num_val)

        total_all = total_all + num_all
        total_train_val = total_train_val + num_train_val
        total_val = total_val + num_val
        total_train = total_train + num_train
        total_test = total_test + num_test

        for f in files_all:
            line = '{}/{} {}\n'.format(class_dir, f, class_idxes[class_dir])

            if f in files_train_val:
                train_val_file.write(line)
                if f in files_val:
                    val_file.write(line)
                else:
                    train_file.write(line)
            else:
                test_file.write(line)

    info_file.write('train: {}/{} \n'.format(total_train, total_all))
    info_file.write('val: {}/{} \n'.format(total_val, total_all))
    info_file.write('trainval: {}/{} \n'.format(total_train_val, total_all))
    info_file.write('test: {}/{} \n'.format(total_test, total_all))
    print 'train: {}/{}'.format(total_train, total_all)
    print 'val: {}/{}'.format(total_val, total_all)
    print 'trainval: {}/{}'.format(total_train_val, total_all)
    print 'test: {}/{}'.format(total_test, total_all)

    train_batch_size = 64
    val_batch_size = 32
    test_iter = total_val / val_batch_size
    max_test_interval = total_train / train_batch_size
    print('')
    print('test_iter = {}'.format(test_iter))
    print('test_interval < {}'.format(max_test_interval))

    train_file.close()
    val_file.close()
    train_val_file.close()
    test_file.close()

    remove_lastline_blank(train_path)
    remove_lastline_blank(val_path)
    remove_lastline_blank(train_val_path)
    remove_lastline_blank(test_path)