pix2pix-mario.py

import skimage
import uuid
import sys
import numpy as np


colors = [
    [92, 148, 252],
    [0,0,0],
    [32,  56, 236]
]

def prepare(img):
    for c in colors:
        img[np.all(img == np.array(c), axis=-1)] = [0,0,0]

    return c

def extract(img, start, width):
    if start + width > img.shape[1]:
        return

    return img[:,start:start+width,:]


def window_ab(img):
    width = height = 224
    stride = 16 * 7
    offset = 0

    while True:
        a = extract(img, offset, width)
        b = extract(img, offset+stride, width)

        if a is None or b is None:
            return

        yield a, b

        offset += 16 * 5

def generate_output_fn():
    return 'train/mario-%s.png' % uuid.uuid4()


def save_sample(a,b):
    out_img = np.zeros([a.shape[0], a.shape[1] + b.shape[1], a.shape[2]], dtype=np.uint8)
    out_img[:,:a.shape[0],:] = a
    out_img[:,a.shape[0]:,:] = b
    skimage.io.imsave(generate_output_fn(), out_img)


def build(gif_filename):
    img = skimage.io.imread(gif_filename)
    standardize_bkg = prepare(img)
    for a, b in window_ab(img):
        save_sample(a, b)


if __name__ == '__main__':
    for gif_filename in sys.argv[1:]:
        build(gif_filename)