import tensorflow as tf
import numpy as np
from PIL import Image
from os import listdir as ls
from sys import argv, exit
from tensorflow.keras.utils import array_to_img as img
from tensorflow.keras import layers as kl

if len(argv) != 3:
    exit(1)

files = ls(argv[1])

for v, fn in enumerate(files):
    i = np.array(Image.open(f"{argv[1]}/{fn}").convert('RGB').resize((300,300))) / 255.0
    print(f"Processing {v+1:03d}/{len(files)}: {fn}")

    for x1 in range(3):
        i1 = tf.image.random_brightness(i, 0.6)
        for x2 in range(3):
            i2 = tf.image.random_saturation(i1, 0.1, 2.0)
            for x3 in range(3):
                i3 = tf.image.random_contrast(i2, 0.2, 1.9)
                img(i3).save(f"{argv[2]}/{fn.rsplit('.', 1)[0]}_{x1}_{x2}_{x3}_1.png")

                i3 = tf.image.flip_left_right(i3)
                img(i3).save(f"{argv[2]}/{fn.rsplit('.', 1)[0]}_{x1}_{x2}_{x3}_2.png")

                i3 = tf.image.flip_up_down(i3)
                img(i3).save(f"{argv[2]}/{fn.rsplit('.', 1)[0]}_{x1}_{x2}_{x3}_3.png")

                i3 = tf.image.flip_left_right(i3)
                img(i3).save(f"{argv[2]}/{fn.rsplit('.', 1)[0]}_{x1}_{x2}_{x3}_4.png")