Generate Pool tiles using Numpy

pool_tiles.py

#!/usr/bin/env python
import numpy as np

def _hex_to_rgb(hex_str):
    """
    Converts hex str to r, g, b
    Asserts that hex_str has a {RRGGBB} format
    """
    r, g, b = hex_str[:2], hex_str[2:4], hex_str[4:]
    rgb = [int(n, base=16) for n in (r,g,b)]
    return rgb[0], rgb[1], rgb[2]


def hex_to_rgb(hex_str):
    """
    Converts hex string to rgb list
    """
    hex_str = hex_str.strip()

    if hex_str[0] == '#':
        hex_str = hex_str[1:]

    if len(hex_str) != 6:
        raise ValueError("Input #{} is not in #RRGGBB format.".format(hex_str))

    r, g, b = hex_str[:2], hex_str[2:4], hex_str[4:]
    rgb = [int(n, base=16) for n in (r,g,b)]
    return rgb


def np_hex_to_rgb(hex_array):
    """
    Converts an array of hex strings to an array of rgb values
    """
    hex_array =  np.char.lstrip(hex_array, '#')

    if hex_array.ndim == 1:
        hex_array = hex_array.reshape(-1, 1)

    nplen = np.vectorize(len)
    if np.any(nplen(hex_array) != 6):
        raise ValueError("The input message should contain strings in #RRGGBB format")

    return np.array(np.vectorize(_hex_to_rgb)(hex_array), dtype=np.uint8).T.reshape(hex_array.shape[0], hex_array.shape[1], 3)

# Pool palette
palette_hex = ['#172636', '#74d7e9', '#27b0cc', '#617685', '#298ccd']
palette_rgb = map(hex_to_rgb, palette_hex)

# Probabilites for each color
chances = [1,70,20,1,8]     # should sum up to 100
chances = np.array(chances, dtype=np.float64) / np.sum(chances) # make sure its normalized and sums to 1

# Display palette
import matplotlib.pyplot as plt
plt.subplot(141)
plt.title('palette')
plt.imshow(np_hex_to_rgb(palette_hex))

# Now Generate random colors in 100x100 shape
tiles = np.random.choice(palette_hex, (100, 100), p=chances)
tiles = np_hex_to_rgb(tiles)

# Display tiles
plt.subplot(142)
plt.title('tiles')
plt.imshow(tiles)

# Zoom tiles
z = 10
zoomed_tiles = np.kron(tiles, np.ones((z,z,1))).astype(np.uint8)

# Plot zoomed tiles
plt.subplot(143)
plt.title('tiles zoomed')
plt.imshow(zoomed_tiles)

# Add white line every n row and line
n = z
pool_tiles = zoomed_tiles.copy()
pool_tiles[:, n::z, :] = [246, 255, 254]
pool_tiles[n::z, :, :] = [246, 255, 254]

# Show Pool tiles
plt.subplot(144)
plt.title('Pool tiles')
plt.imshow(pool_tiles)


plt.show()


# Save array as image file
import scipy.misc
scipy.misc.imsave('pool_titles.png',  pool_tiles)
print 'Saved file as pool_tiles.png'