Script to generate images using a markov chain trained on hilbert curve mapped images

Hilbert-Markov.py

import markovify
import hilbert
import os
import os.path
import numpy as np
from PIL import Image


def rgb2hex(rgb):
    hex_result = "".join([str(format(val, '02x')).upper() for val in rgb])
    return f"{hex_result}"


def hex2rgb(color):
    return list(int(color[i:i + 2], 16) for i in (0, 2, 4))


def get_images(dir):
    valid_images = [".jpg", ".gif", ".png", ".tga"]
    for f in os.listdir(dir):
        ext = os.path.splitext(f)[1]
        if ext.lower() not in valid_images:
            continue
        yield Image.open(os.path.join(dir, f))


def d2_to_arr(m, matrix):
    result = [None] * (m ** 2)
    for i in range(m):
        for j in range(m):
            result[hilbert.xy2d(m, j, i)] = matrix[i][j]

    return result


def arr_to_d2(m, arr):
    result = [[None] * m for _ in range(m)]
    for i in range(len(arr)):
        x, y = hilbert.d2xy(m, i)
        result[y][x] = arr[i]

    return result


IMG_DIR = './assorted_2'
SAVE_DIR = './assorted_2/generated'
# the hilbert curve mapping only works on squares whose sides are a power of 2
IMG_SIZE = 512
STATE_SIZE = 4
GENERATE_N = 100

corpus = []
for image in get_images(IMG_DIR):
    image = image.convert("RGB")
    picture = list(image.getdata())
    # take the 1d array from PIL and turn it into a 2d array for d2_to_arr 
    picture = [[list(picture[(i * IMG_SIZE) + j]) for j in range(IMG_SIZE)] for i in range(IMG_SIZE)]
    picture = [[rgb2hex(j) for j in i] for i in picture]
    # turns the image into a "sentence" using the hilbert curve
    hilbert_picture = d2_to_arr(IMG_SIZE, picture)
    corpus.append(hilbert_picture)
    print("Image Processed")

# trains the markov chain on our group of "sentences"
chain = markovify.Chain(corpus, STATE_SIZE)
print("Chain Trained")

generated = 0
while generated < GENERATE_N:
    hilbert_picture = list(chain.gen())
    # hilbert curve mapping only works on squares whose side is a power of 2
    # you could do something like have various size images outputed that just get truncated down to the closest (smaller) power of 2
    # or go to the closest larger power of 2 and fill in the missing pixels with black or transparency or whatever
    # but i just wanted images of the same size i trained on so i make sure the "sentence" is long enough and then truncate
    if len(hilbert_picture) >= (IMG_SIZE ** 2):
        del hilbert_picture[(IMG_SIZE ** 2):]
        for i, item in enumerate(hilbert_picture):
            hilbert_picture[i] = hex2rgb(item)
        # reverses the hilbert curve mapping and save the image
        picture = arr_to_d2(IMG_SIZE, hilbert_picture)
        image = Image.fromarray(np.asarray(picture).astype('uint8'))
        # image.show()
        image.save(f"{SAVE_DIR}/{generated}.png")
        generated += 1