O2020/make_dataset.py

import os
import sys
import yaml
import json
from utils.dataset_config import dataset_config
from utils.misc import seed_everything


'''
Main function of this script. It is explained in more detail in the tutorial.
'''
def run(parameters, directories):

    #Retrieving directories information
    data_path = directories['data_path']
    train_test_path = directories['train_test_path']
    csv_path = directories['csv_path']
    train_dir = os.path.join(train_test_path, "train")
    val_dir = os.path.join(train_test_path, "val")
    test_dir = os.path.join(train_test_path, "test")

    #Retrieving parameters
    seed = parameters['seed']
    val_size = parameters['val_size']
    test_size = parameters['test_size']

    #Setting random processes seed to assure reproducibility
    seed_everything(seed)

    #Gathering dataset's images information from their folder
    dataset = dataset_config(dir=data_path)

    #Getting train and test images' ids
    if os.path.isdir(train_dir) and os.path.isdir(val_dir) and os.path.isdir(test_dir):
        #In the case that the dataset has already been split
        train_ids, val_ids, test_ids = dataset.load(train_test_path)
    else:
        #In the case that the dataset has not yet been split
        train_ids, val_ids, test_ids = dataset.train_split(path=train_test_path, csv_path=csv_path, val_size=val_size, test_size=test_size)

    ids_dict = {'train_ids': train_ids, 'val_ids': val_ids, 'test_ids': test_ids}

    with open('ids_dict.json', 'w') as fp:
        json.dump(ids_dict, fp)

if __name__ == "__main__":

    #Attempting to read dir.yaml file that contains directories information
    try:
        with open(os.path.join("configs","dir.yaml"), "r") as fp:
            directories = yaml.load(fp, Loader=yaml.FullLoader)
    except:
        sys.exit('File dir.yaml does not exist in the current folder')

    #Attempting to read parameters.yaml file that contains parameters information
    try:
        with open(os.path.join("configs","parameters.yaml"), "r") as fp:
            parameters = yaml.load(fp, Loader=yaml.FullLoader)
    except:
        sys.exit('File parameters.yaml does not exist in the current folder')

    run(parameters, directories)