import sys
sys.path.append('..')

import glob
from typing import NoReturn
from multiprocessing import cpu_count
from sklearn.model_selection import train_test_split
from src.utilities.tfrecords_save import multi_create_tfrecords_in_shards


def main(
    num_files_to_create: int = 96,
    num_workeers: int = cpu_count() - 1
    ) -> NoReturn:

    num_files_to_create = 96
    num_workers = cpu_count()-1

    
    filenames = {
        'train': None,
        'test' : None,
        'validation' : None
    }

    image_files = glob.glob(f'/data/spec/*/*')
    train_filenames, validation_test_filenames = train_test_split(
        image_files,
        train_size = 0.8,
        test_size = 0.2,
        random_state = 24601
    )

    validation_filenames, test_filenames = train_test_split(
        validation_test_filenames,
        train_size = 0.5,
        test_size = 0.5,
        random_state = 24601
    )


    num_test_samples, test_tfrecords = multi_create_tfrecords_in_shards(
        data_filenames = test_filenames,
        dataset_name = 'test',
        num_files_to_create = num_files_to_create,
        num_workers = num_workers
    )


    num_validation_samples, validation_tfrecords = multi_create_tfrecords_in_shards(
        data_filenames = validation_filenames,
        dataset_name = 'validation',
        num_files_to_create = num_files_to_create,
        num_workers = num_workers
    )


    num_train_samples, train_tfrecords = multi_create_tfrecords_in_shards(
        data_filenames = train_filenames,
        dataset_name = 'train',
        num_files_to_create = num_files_to_create,
        num_workers = num_workers
    )
    
    
    filenames['train'] = train_tfrecords
    filenames['test'] = test_tfrecords
    filenames['validation'] = validation_tfrecords
    
    return filenames


if __name__ == "__main__":
    main()

import sys
sys.path.append('..')

import glob
from typing import NoReturn
from multiprocessing import cpu_count
from sklearn.model_selection import train_test_split
from src.utilities.tfrecords_save import multi_create_tfrecords_in_shards


def main(
    num_files_to_create: int = 96,
    num_workeers: int = cpu_count() - 1
    ) -> NoReturn:

    num_files_to_create = 96
    num_workers = cpu_count()-1

    
    filenames = {
        'train': None,
        'test' : None,
        'validation' : None
    }

    image_files = glob.glob(f'/data/spec/*/*')
    train_filenames, validation_test_filenames = train_test_split(
        image_files,
        train_size = 0.8,
        test_size = 0.2,
        random_state = 24601
    )

    validation_filenames, test_filenames = train_test_split(
        validation_test_filenames,
        train_size = 0.5,
        test_size = 0.5,
        random_state = 24601
    )


    num_test_samples, test_tfrecords = multi_create_tfrecords_in_shards(
        data_filenames = test_filenames,
        dataset_name = 'test',
        num_files_to_create = num_files_to_create,
        num_workers = num_workers
    )


    num_validation_samples, validation_tfrecords = multi_create_tfrecords_in_shards(
        data_filenames = validation_filenames,
        dataset_name = 'validation',
        num_files_to_create = num_files_to_create,
        num_workers = num_workers
    )


    num_train_samples, train_tfrecords = multi_create_tfrecords_in_shards(
        data_filenames = train_filenames,
        dataset_name = 'train',
        num_files_to_create = num_files_to_create,
        num_workers = num_workers
    )
    
    
    filenames['train'] = train_tfrecords
    filenames['test'] = test_tfrecords
    filenames['validation'] = validation_tfrecords
    
    return filenames


if __name__ == "__main__":
    main()
