<commit_before>import os
import subprocess

import tqdm

import numpy as np

from keras.preprocessing.image import ImageDataGenerator


aug_rounds = 4


# TODO: replace this listdir with a mapping tbl/json
CATS = sorted(os.listdir('data/raw/train'))


def cat_from_int(cat_int):
    return CATS[cat_int]


def gen_XY_from_dir(arr_dir):
    arr_files = sorted(os.listdir(arr_dir))
    arr_names = list(filter(lambda x: r'-img-' in x, arr_files))
    lab_names = list(filter(lambda x: r'-lab-' in x, arr_files))

    assert len(arr_names) == len(lab_names), '# labels != images'

    for arr_name, lab_name in zip(arr_names, lab_names):
        X = np.load(os.path.join(arr_dir, arr_name))
        Y = np.load(os.path.join(arr_dir, lab_name))
        yield X, Y


def augment_XY(x, y, aug_rounds):
    auggen = ImageDataGenerator(rotation_range=10,
                                width_shift_range=0.1,
                                height_shift_range=0.1,
                                horizontal_flip=True)

    X_aug, Y_aug = next(auggen.flow(np.tile(x[np.newaxis],
                                            (aug_rounds, 1, 1, 1)),
                                    np.tile(y[np.newaxis],
                                            (aug_rounds, 1)),
                                    batch_size=aug_rounds))

    for x_aug, y_aug in zip(X_aug, Y_aug):
        yield x_aug, y_aug


def arrs_to_aug(arr_dir, aug_dir):

    subprocess.call(['mkdir', '-p', aug_dir])

    for img_idx, (x, y) in tqdm.tqdm(enumerate(gen_XY_from_dir(arr_dir))):
        for aug_idx, (x_aug, y_aug) in enumerate(augment_XY(x, y, aug_rounds)):
            cat_idx = np.argmax(y_aug)
            cat = cat_from_int(cat_idx)
            img_name = '{:04d}-{:02d}-img-{}-{}'.format(img_idx, aug_idx,
                                                        cat, cat_idx)
            lab_name = '{:04d}-{:02d}-lab-{}-{}'.format(img_idx, aug_idx,
                                                        cat, cat_idx)
            aug_path = os.path.join(aug_dir, img_name)
            lab_path = os.path.join(aug_dir, lab_name)
            np.save(aug_path, x_aug)
            np.save(lab_path, y_aug)


if __name__ == '__main__':
    arrs_to_aug('data/arr/train', 'data/aug/train')
    arrs_to_aug('data/arr/test', 'data/aug/test')<commit_msg>Remove unused augmentation on test set<commit_after>import os
import subprocess
import tqdm
import numpy as np
from keras.preprocessing.image import ImageDataGenerator

aug_rounds = 4
# TODO: replace this listdir with a mapping tbl/json
CATS = sorted(os.listdir('data/raw/train'))


def cat_from_int(cat_int):
    return CATS[cat_int]


def gen_XY_from_dir(arr_dir):
    arr_files = sorted(os.listdir(arr_dir))
    arr_names = list(filter(lambda x: r'-img-' in x, arr_files))
    lab_names = list(filter(lambda x: r'-lab-' in x, arr_files))

    assert len(arr_names) == len(lab_names), '# labels != images'

    for arr_name, lab_name in zip(arr_names, lab_names):
        X = np.load(os.path.join(arr_dir, arr_name))
        Y = np.load(os.path.join(arr_dir, lab_name))
        yield X, Y


def augment_XY(x, y, aug_rounds):
    auggen = ImageDataGenerator(rotation_range=10,
                                width_shift_range=0.1,
                                height_shift_range=0.1,
                                horizontal_flip=True)

    X_aug, Y_aug = next(auggen.flow(np.tile(x[np.newaxis],
                                            (aug_rounds, 1, 1, 1)),
                                    np.tile(y[np.newaxis],
                                            (aug_rounds, 1)),
                                    batch_size=aug_rounds))

    for x_aug, y_aug in zip(X_aug, Y_aug):
        yield x_aug, y_aug


def arrs_to_aug(arr_dir, aug_dir):

    subprocess.call(['mkdir', '-p', aug_dir])

    for img_idx, (x, y) in tqdm.tqdm(enumerate(gen_XY_from_dir(arr_dir))):
        for aug_idx, (x_aug, y_aug) in enumerate(augment_XY(x, y, aug_rounds)):
            cat_idx = np.argmax(y_aug)
            cat = cat_from_int(cat_idx)
            img_name = '{:04d}-{:02d}-img-{}-{}'.format(img_idx, aug_idx,
                                                        cat, cat_idx)
            lab_name = '{:04d}-{:02d}-lab-{}-{}'.format(img_idx, aug_idx,
                                                        cat, cat_idx)
            aug_path = os.path.join(aug_dir, img_name)
            lab_path = os.path.join(aug_dir, lab_name)
            np.save(aug_path, x_aug)
            np.save(lab_path, y_aug)


if __name__ == '__main__':
    arrs_to_aug('data/arr/train', 'data/aug/train')
{'fraction_non_alphanumeric': 0.12102186620480623, 'fraction_numerical': 0.008226888936999351, 'mean_word_length': 1.9129886506935687, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 20, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6683183', 'n_tokens_mistral': 1621, 'n_tokens_neox': 1529, 'n_words': 298}