Image classification models

.
└── data_root/
    ├── images/
    ├── labels.csv
    └── class_to_idx.json  # Optional

import asyncio
import aiofiles
from io import BytesIO
from PIL import Image
import pickle
from pathlib import Path
import pandas as pd
import random


# Change to false if only the labels.csv file needs to be processed
SAVE_IMAGE = True

# The async code will open too many files at one time. Let's limit this
num_of_max_files_open = 200

data_dir = Path('./data')
data_dir.mkdir(exist_ok=True)


def unpickle(file):
    with open(file, 'rb') as fo:
        data = pickle.load(fo, encoding='bytes')
    return data


def load_subset(subset):
    data = unpickle(f'./cifar-100-python/{subset}')

    filenames = data[b'filenames']
    labels = data[b'fine_labels']
    images = data[b'data']

    assert len(labels) == len(images)
    assert len(filenames) == len(images)

    return filenames, labels, images


async def save_image(path: str, image: memoryview) -> None:
    async with aiofiles.open(path, "wb") as file:
        await file.write(image)


async def write_image(filename, label, image, subset, row, sem):
    subset_dir = data_dir / subset
    filepath = subset_dir / filename.decode()

    async with sem:
        # Each row of the array stores a 32x32 colour image. The first 1024 entries contain the red channel values, the
        # next 1024 the green, and the final 1024 the blue. The image is stored in row-major order, so that the first 32
        # entries of the array are the red channel values of the first row of the image.
        if SAVE_IMAGE:
            image = image.reshape(3, 32, 32).transpose(1, 2, 0)
            img = Image.fromarray(image)
            buffer = BytesIO()
            img.save(buffer, format='png')
            await save_image(filepath, buffer.getbuffer())

    if row % 100 == 0:
        print(f"{row:05d}", flush=True)

    if subset == 'test':
        # we use the ``test`` set as the ``validation`` set in this example
        subset = 'validation'

    return [str(filepath.relative_to(data_dir)), str(label), subset, str('')]


async def process_subset(subset):

    subset_dir = data_dir / subset
    subset_dir.mkdir(exist_ok=True)

    tasks = []
    sem = asyncio.Semaphore(num_of_max_files_open)
    filenames, labels, images = load_subset(subset)

    for row, sample in enumerate(zip(filenames, labels, images)):
        tasks.append(asyncio.ensure_future(write_image(*sample, subset=subset, row=row, sem=sem)))

    results = await asyncio.gather(*tasks)
    df = pd.DataFrame(results, columns=["image_path", "label", "subset", "metadata"])
    return df


async def main():

    print("Processing training images")
    train_df = await process_subset('train')
    print("Processing test images")
    test_df = await process_subset('test')
    df = pd.concat([train_df, test_df])
    df.to_csv(data_dir / 'labels.csv', index=False)


asyncio.run(main())

.
└── data_root/
    ├── images/
    └── predictions.csv