Boilerplate



In [ ]:

    
import copy
import os

import matplotlib.pyplot as plt
import numpy as np
from scipy import ndimage
import tensorflow as tf
import tensorflow_hub as hub

import spiral.agents.default as default_agent
import spiral.agents.utils as agent_utils
from spiral.environments import fluid
from spiral.environments import libmypaint


nest = tf.contrib.framework.nest

# Disable TensorFlow debug output.
tf.logging.set_verbosity(tf.logging.ERROR)

Sample from the model

This section demonstrates the most basic usage of the package, i.e., sampling from a pre-trained model.



In [ ]:

    
# The path to libmypaint brushes.
BRUSHES_BASEDIR = os.path.join(os.getcwd(), "..", "third_party/mypaint-brushes-1.3.0")
BRUSHES_BASEDIR = os.path.abspath(BRUSHES_BASEDIR)
# The path to a TF-Hub module.
MODULE_PATH = "https://tfhub.dev/deepmind/spiral/default-wgangp-celebahq64-gen-19steps/agent4/1"

Create the environment

First, we need to create the environment.



In [ ]:

    
env_settings = dict(
    episode_length=20,                 # Number of frames in each episode.
    canvas_width=64,                   # The width of the canvas in pixels.
    grid_width=32,                     # The width of the action grid.
    brush_type="classic/dry_brush",    # The type of the brush.
    brush_sizes=[1, 2, 4, 8, 12, 24],  # The sizes of the brush to use.
    use_color=True,                    # Color or black & white output?
    use_pressure=True,                 # Use pressure parameter of the brush?
    use_alpha=False,                   # Drop or keep the alpha channel of the canvas?
    background="white",                # Background could either be "white" or "transparent".
    brushes_basedir=BRUSHES_BASEDIR,   # The location of libmypaint brushes.
)
env = libmypaint.LibMyPaint(**env_settings)

Create the agent

We provide a convenience function get_module_wrappers that returns two python functions implementing the agent. The first one, initial_state, is used to get the initial state of the agent (specifically, the state of the LSTM). The second, step, takes an observation from the environment and performs a single agent step.



In [ ]:

    
initial_state, step = agent_utils.get_module_wrappers(MODULE_PATH)

Run sampling



In [ ]:

    
state = initial_state()

Running the following cell will return a sample from the model. You can execute it multiple times until you get the one that you like. In addition to showing the final state of the canvas we also record all the agent's actions (note actions variable) needed to reproduce it.



In [ ]:

    
noise_sample = np.random.normal(size=(10,)).astype(np.float32)

time_step = env.reset()
actions = []
for t in range(19):
    time_step.observation["noise_sample"] = noise_sample
    action, state = step(time_step.step_type, time_step.observation, state)
    time_step = env.step(action)
    actions.append(action)
    
plt.close("all")
plt.figure(figsize=(5, 5))
plt.imshow(time_step.observation["canvas"], interpolation="nearest")

Manipulate the sample

Let's now do something the obtained sample. Since we have the corresponding sequence of actions we can re-render the image in higher resolution. To that end, we will need to create one more environment with modified settings.

Create a high-resolution version of the environment



In [ ]:

    
# Let's make the canvas 8x8 times bigger.
SCALE_FACTOR = 8

# Patch the environments setting for higher resolution.
hires_env_settings = copy.deepcopy(env_settings)
hires_env_settings["canvas_width"] *= SCALE_FACTOR
hires_env_settings["brush_sizes"] = [
    s * SCALE_FACTOR for s in hires_env_settings["brush_sizes"]]

env = libmypaint.LibMyPaint(**hires_env_settings)

Execute pre-recorded actions



In [ ]:

    
env.reset()
for t in range(19):
    time_step = env.step(actions[t])
    
plt.close("all")
plt.figure(figsize=(5, 5))
plt.imshow(time_step.observation["canvas"], interpolation="nearest")

Change the thickness of the brush strokes

In addition to changing the resolution of the images, let's in introduce some more subtle structural changes. We could, for example, change the thickness of all the strokes.



In [ ]:

    
modified_actions = copy.deepcopy(actions)
for action in modified_actions:
    action["size"] = np.array(0, dtype=np.int32)
    action["pressure"] = np.array(2, dtype=np.int32)

env.reset()
for t in range(19):
    time_step = env.step(modified_actions[t])
    
plt.close("all")
plt.figure(figsize=(5, 5))
plt.imshow(time_step.observation["canvas"], interpolation="nearest")

Change the brush type

Finally, let's re-render the image above using a different brush type.



In [ ]:

    
# Patch the environments setting for a different brush.
pen_env_settings = copy.deepcopy(hires_env_settings)
pen_env_settings["brush_type"] = "classic/pen"

env = libmypaint.LibMyPaint(**pen_env_settings)



In [ ]:

    
env.reset()
for t in range(19):
    time_step = env.step(modified_actions[t])
    
plt.close("all")
plt.figure(figsize=(5, 5))
plt.imshow(time_step.observation["canvas"], interpolation="nearest")

Fluid Paint environment demonstration

Fluid Paint environment works almost exactly like libmypaint.LibMyPaint. Below, we show how to obtain samples from the model trained in this environment.



In [ ]:

    
# The path to the shaders.
SHADERS_BASEDIR = os.path.join(os.getcwd(), "..", "third_party/paint/shaders")
SHADERS_BASEDIR = os.path.abspath(SHADERS_BASEDIR)
# The path to a TF-Hub module.
MODULE_PATH = "https://tfhub.dev/deepmind/spiral/default-fluid-gansn-celebahq64-gen-19steps/1"



In [ ]:

    
env_settings = dict(
    episode_length=20,                               # Number of frames in each episode.
    canvas_width=256,                                # The width of the canvas in pixels.
    grid_width=32,                                   # The width of the action grid.
    brush_sizes=[2.5, 5.0, 10.0, 20.0, 40.0, 80.0],  # The sizes of the brush to use.
    shaders_basedir=SHADERS_BASEDIR,                 # The location of shaders.
)
env = fluid.FluidPaint(**env_settings)



In [ ]:

    
initial_state, step = agent_utils.get_module_wrappers(MODULE_PATH)



In [ ]:

    
state = initial_state()



In [ ]:

    
noise_sample = np.random.normal(size=(10,)).astype(np.float32)

time_step = env.reset()
actions = []
for t in range(19):
    time_step.observation["noise_sample"] = noise_sample
    # The environment uses 256x256 canvas but the agent requires 64x64 input.
    ratio = 64 / 256
    time_step.observation["canvas"] = ndimage.zoom(
        time_step.observation["canvas"], [ratio, ratio, 1], order=1)
    action, state = step(time_step.step_type, time_step.observation, state)
    time_step = env.step(action)
    actions.append(action)
    
plt.close("all")
plt.figure(figsize=(5, 5))
plt.imshow(time_step.observation["canvas"], interpolation="nearest")