training_generator_example.py

#!/usr/bin/env python3

import os
from PIL import Image, ImageFilter

from enum import Enum
import random
import numpy as np

import tensorflow as tf
import datetime
import math
import time


def is_png(f):
    return f.lower().endswith('.png')


def is_jpeg(f):
    return f.lower().endswith('.jpg') or f.lower().endswith('.jpeg')


VAL_INPUT_FOLDER = 'val_input'
VAL_OUTPUT_FOLDER = 'val_output'
TRAIN_INPUT_FOLDER = 'train_input'
TRAIN_OUTPUT_FOLDER = 'train_output'

VAL_INPUT_LIST = [
    f for f in (os.listdir(VAL_INPUT_FOLDER)) if (is_png(f) or is_jpeg(f))
]

VAL_OUTPUT_LIST = [
    f for f in (os.listdir(VAL_OUTPUT_FOLDER)) if (is_png(f) or is_jpeg(f))
]

TRAIN_INPUT_LIST = [
    f for f in (os.listdir(TRAIN_INPUT_FOLDER)) if (is_png(f) or is_jpeg(f))
]

TRAIN_OUTPUT_LIST = [
    f for f in (os.listdir(TRAIN_OUTPUT_FOLDER)) if (is_png(f) or is_jpeg(f))
]

VAL_INPUT_LIST.sort()
VAL_OUTPUT_LIST.sort()
TRAIN_INPUT_LIST.sort()
TRAIN_OUTPUT_LIST.sort()

assert VAL_INPUT_LIST == VAL_OUTPUT_LIST
assert TRAIN_INPUT_LIST == TRAIN_OUTPUT_LIST

#####################################################################

MAX_BLOCK_WIDTH = 160
MAX_BLOCK_HEIGHT = 160


class DataPurpose(Enum):
    TRAINING = 1
    VALIDATION = 2


def get_random_image_pair(purpose: DataPurpose):
    if purpose == DataPurpose.VALIDATION:
        settings = (VAL_INPUT_FOLDER, VAL_OUTPUT_FOLDER, VAL_INPUT_LIST)
    else:
        settings = (TRAIN_INPUT_FOLDER, TRAIN_OUTPUT_FOLDER, TRAIN_INPUT_LIST)

    name = random.choice(settings[2])

    input = Image.open(os.path.join(settings[0], name))
    output = Image.open(os.path.join(settings[1], name))

    assert input.size[0] * 2 == output.size[0]
    assert input.size[1] * 2 == output.size[1]

    return (input, output)


def random_image_block_pair(image_pair):
    input, output = image_pair

    assert input.size[0] * 2 == output.size[0]
    assert input.size[1] * 2 == output.size[1]

    bw = min(MAX_BLOCK_WIDTH, input.size[0])
    bh = min(MAX_BLOCK_HEIGHT, input.size[1])

    x = random.randint(0, input.size[0] - bw)
    y = random.randint(0, input.size[1] - bh)

    x2 = x + bw
    y2 = y + bh

    big_x = 2 * x
    big_y = 2 * y
    big_x2 = big_x + 2 * bw
    big_y2 = big_y + 2 * bh

    input = input.crop((x, y, x2, y2))
    output = output.crop((big_x, big_y, big_x2, big_y2))

    assert input.size[0] == bw, "{} != {}".format(input.size[0], bw)
    assert input.size[1] == bh, "{} != {}".format(input.size[1], bh)

    assert input.size[0] * 2 == output.size[0]
    assert input.size[1] * 2 == output.size[1]

    return input, output


def training_generator():
    while (True):

        images_per_step = 10
        blocks_per_image = 5

        batch_size = images_per_step * blocks_per_image

        bw = 2 * MAX_BLOCK_WIDTH
        bh = 2 * MAX_BLOCK_HEIGHT
        inputs, outputs = np.zeros((batch_size, MAX_BLOCK_HEIGHT, MAX_BLOCK_WIDTH, 4), dtype=np.float32), \
                          np.zeros((batch_size, bh, bw, 2), dtype=np.float32)

        image_index = 0

        while image_index < images_per_step:

            image_pair = get_random_image_pair(DataPurpose.TRAINING)
            
            for block_index in range(blocks_per_image):
                pair = random_image_block_pair(image_pair)
                
                # Some color space or other converting here...
                input_image_array = TODO
                output_image_array = TODO
                
                inputs[blocks_per_image * image_index + block_index] = input_image_array
                outputs[blocks_per_image * image_index + block_index] = output_image_array
            
            image_index += 1
            
        yield inputs, outputs