Converter for Cifar100 pickle data to png

cifar100_to_png.py

from PIL import Image
import numpy as np
import pickle
import argparse
import os

metadata_dict = None


def unpickle(file_path):
    result = None
    with open(file_path, 'rb') as f:
        result = pickle.load(f, encoding='latin-1')
    return result


def convert_image(dataset, index):
    data = np.rollaxis(dataset['data'][index].reshape((3, 32, 32)), 0, 3)
    return Image.fromarray(data)


def init_metadata_dict(metadata_path):
    global metadata_dict
    if metadata_dict == None:
        with open(metadata_path, 'rb') as f:
            metadata_dict = unpickle(metadata_path)

    return metadata_dict


def get_fine_label(dataset, index, use_humanistic_label=False):
    if use_humanistic_label:
        return metadata_dict['fine_label_names'][dataset['fine_labels'][index]]
    else:
        return dataset['fine_labels'][index]


def get_coarse_label(dataset, index, use_humanistic_label=False):
    if use_humanistic_label:
        return metadata_dict['coarse_label_names'][dataset['coarse_labels'][index]]
    else:
        return dataset['coarse_labels'][index]


def main():
    parser = argparse.ArgumentParser(
        description='Image converter for Chainer100 dataset')
    parser.add_argument('--dataset_path', '-d',
                        default='cifar-100-python/train')
    parser.add_argument('--output_path', '-o', default='./images')
    parser.add_argument('--metadata_path', '-m',
                        default='cifar-100-python/meta')
    parser.add_argument('--convert_from', '-f', type=int, default=0)
    parser.add_argument('--convert_to', '-t', type=int, default=1000)
    parser.add_argument('--use_humanistic_label', type=bool, default=False)
    args = parser.parse_args()

    print(f'dataset_path: {args.dataset_path}')
    print(f'output_path: {args.output_path}')

    if args.use_humanistic_label:
        init_metadata_dict(args.metadata_path)
    os.makedirs(args.output_path, exist_ok=True)

    dataset = unpickle(args.dataset_path)
    print(dataset['data'].shape)
    for index in range(args.convert_from, args.convert_to):
        image = convert_image(dataset, index)
        coarse_label = get_coarse_label(
            dataset, index, args.use_humanistic_label)
        fine_label = get_fine_label(dataset, index, args.use_humanistic_label)
        image_save_directory = f'{args.output_path}/{coarse_label}/{fine_label}'
        os.makedirs(image_save_directory, exist_ok=True)
        image.save(f'{image_save_directory}/{index}.png')

    print('done')


if __name__ == '__main__':
    main()