Skip to content

Commit

Permalink
update docstrings, refactor variables
Browse files Browse the repository at this point in the history
  • Loading branch information
@krachbumm3nte
krachbumm3nte committed May 20, 2022
1 parent 4194ee7 commit 93dfaf2
Showing 1 changed file with 70 additions and 62 deletions.
132 changes: 70 additions & 62 deletions pynest/examples/pong/generate_gif.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -23,12 +23,11 @@
----------------------------------------------------------------
All simulations store data about both networks and the game in .pkl files.
This script reads these files and generates image snapshots at different
times during the simulation. These are subsequently aggregated into a GIF
times during the simulation. These are subsequently aggregated into a GIF.
:Authors: J Gille, T Wunderlich, Electronic Vision(s)
"""


import numpy as np
from PIL import Image, ImageFont, ImageDraw
import pickle
Expand All @@ -38,96 +37,99 @@
import io
import sys
from copy import copy
import pong
from pong import GameOfPong as Pong, Paddle
import imageio
from glob import glob
from matplotlib import font_manager


# base colors used in the plot
background_color = (220, 220, 220) # grey
background_hex = "#dcdcdc" # hex value of the same color, required for matplotlib
background_hex = "#dcdcdc" # hex value of grey, required for matplotlib
black = np.array((0, 0, 0))
white = np.array((255, 255, 255))
left_color = np.array((204, 0, 153)) # purple
right_color = np.array((255, 128, 0)) # orange

# final output image size
IMAGE_SIZE = np.array([400, 320])
# original size of the playing field inside the simulation
FIELD_SIZE = np.array([pong.GameOfPong.x_grid, pong.GameOfPong.y_grid])
# grid size of the playing field inside the simulation
GAME_GRID = np.array([Pong.x_grid, Pong.y_grid])

FIELD_SCALE = 8
# Field size (in px) in the final image
FIELD_SIZE_SCALED = FIELD_SIZE * FIELD_SCALE
GAME_GRID_SCALED = GAME_GRID * FIELD_SCALE

# dimensions of game objects in px
BALL_RAD = 2
PADDLE_LEN = int(0.1*FIELD_SIZE_SCALED[1])
PADDLE_LEN = int(0.5*Paddle.length*Pong.y_length*GAME_GRID_SCALED[1])
PADDLE_WID = 6

# Add margins left and right to the image representing the playing field
# Add margins left and right to the playing field
FIELD_PADDING = PADDLE_WID * 2
FIELD_IMAGE_SIZE = copy(FIELD_SIZE_SCALED)
FIELD_IMAGE_SIZE[0] += 2*FIELD_PADDING
FIELD_SIZE = copy(GAME_GRID_SCALED)
FIELD_SIZE[0] += 2*FIELD_PADDING

HEATMAP_SCALE = 4
# weight matrix heatmap size (in px) in the final image
HEATMAP_SIZE = np.array([FIELD_SIZE[1], FIELD_SIZE[1]]) * HEATMAP_SCALE
HEATMAP_SIZE = np.array([GAME_GRID[1], GAME_GRID[1]]) * HEATMAP_SCALE


DEFAULT_SPEED = 4
PLOT_INTERVAL = 10

out_file = "pong_sim.gif"
keep_temps = False


def scale_coordinates(coordinates: np.array):
"""Scale an (x,y) coordinate tuple from simulation scale to a pixel
coordinate in the output image
coordinate in the output image.
Args:
pos (float, float): input coordinates to be scaled
pos (float, float): input coordinates to be scaled.
Returns:
(int, int): output coordinates in pixel value
(int, int): output coordinates in the image (in px).
"""
coordinates[:, 0] = coordinates[:, 0] * \
FIELD_SIZE_SCALED[0] / pong.GameOfPong.x_length + FIELD_PADDING
coordinates[:, 1] = coordinates[:, 1] * \
FIELD_SIZE_SCALED[1] / pong.GameOfPong.y_length
coordinates[:, 0] = coordinates[:, 0] * GAME_GRID_SCALED[0] \
/ Pong.x_length + FIELD_PADDING
coordinates[:, 1] = coordinates[:, 1] * GAME_GRID_SCALED[1] \
/ Pong.y_length
return coordinates.astype(int)


def grayscale_to_heatmap(in_image, min_val, max_val, base_color):
"""transform a grayscale image to an RGB heat map. Heatmap will color small
values in base_color and high values white
values in base_color and high values in white.
Args:
in_image (numpy.array): 2D numpy.array to be transformed
min_val (float): smallest value across the entire dataset - colored in
base_color in the output
max_val (float): largest value across the entire dataset - colored
white (255, 255, 255)
in_image (numpy.array): 2D numpy.array to be transformed.
min_val (float): smallest value across the entire image - colored in
base_color in the output.
max_val (float): largest value across the entire image - colored
white in the output.
base_color (numpy.array): numpy.array of shape (3,) representing the
base color of the heatmap in RGB space
base color of the heatmap in RGB.
Returns:
numpy.array: transformed input array with an added 3rd dimension of
length three representing RGB values
length three representing RGB values.
"""

x_len, y_len = in_image.shape
out_image = np.zeros((x_len, y_len, 3), dtype=np.uint8)

# add miniscule offsett to avoid zero division error for uniform matrix
# add miniscule offsett to avoid zero division error for uniform matrix.
span = max_val - min_val + 0.0001
for i in range(x_len):
for j in range(y_len):
val = (in_image[i, j] - min_val) / span
out_image[i, j, :] = base_color + (white-base_color) * val
color_scaled = (in_image[i, j] - min_val) / span
out_image[i, j, :] = base_color + (white-base_color) * color_scaled

return out_image


if __name__ == "__main__":
# plotting parameters
keep_temps = False
out_file = "pong_sim.gif"

if len(sys.argv) != 2:
print("This programm takes exactly one argument - the location of the "
Expand All @@ -146,8 +148,10 @@ def grayscale_to_heatmap(in_image, min_val, max_val, base_color):
else:
os.mkdir(temp_dir)

# retrieve fonts independent of the platform used.
font_prop = font_manager.FontProperties(family="sans")
font_name = font_manager.findfont(font_prop)

font_large = ImageFont.truetype(font_name, 26, encoding="unic")
font_medium = ImageFont.truetype(font_name, 18, encoding="unic")

Expand All @@ -158,63 +162,64 @@ def grayscale_to_heatmap(in_image, min_val, max_val, base_color):
ball_positions = scale_coordinates(np.array(game_data["ball_pos"]))
ball_positions[:, 0] -= BALL_RAD
l_paddle_positions = scale_coordinates(np.array(game_data["left_paddle"]))
# move left paddle outwards for symmetry
# move left paddle outwards for symmetry.
l_paddle_positions[:, 0] -= PADDLE_WID
r_paddle_positions = scale_coordinates(np.array(game_data["right_paddle"]))

score = np.array(game_data["score"]).astype(int)

with gzip.open(os.path.join(input_folder, "data_left.pkl.gz"), 'r') as f:
data = pickle.load(f)
rewards_left = data["rewards"]
weights_left = data["weights"]
p1_name = data["network_type"]
name_left = data["network_type"]

with gzip.open(os.path.join(input_folder, "data_right.pkl.gz"), 'r') as f:
data = pickle.load(f)
rewards_right = data["rewards"]
weights_right = data["weights"]
p2_name = data["network_type"]
name_right = data["network_type"]

# extract lowest and highest weights for both players to scale the heatmaps
# extract lowest and highest weights for both players to scale the heatmaps.
min_r, max_r = np.min(weights_right), np.max(weights_right)
min_l, max_l = np.min(weights_left), np.max(weights_left)

# average rewards at every iteration over all neurons
# average rewards at every iteration over all neurons.
rewards_left = [np.mean(x) for x in rewards_left]
rewards_right = [np.mean(x) for x in rewards_right]

sim_iterations = score.shape[0]
n_iterations = score.shape[0]
i = 0
output_speed = DEFAULT_SPEED

print(f"setup complete, generating images to {temp_dir}...")
while i < sim_iterations:
while i < n_iterations:

background = Image.new("RGB", tuple(IMAGE_SIZE), background_color)
draw = ImageDraw.Draw(background)

# create an empty array for the playing field.
playing_field = np.zeros(
(FIELD_IMAGE_SIZE[0], FIELD_IMAGE_SIZE[1], 3), dtype=np.uint8)
(FIELD_SIZE[0], FIELD_SIZE[1], 3), dtype=np.uint8)

# draw Ball in white
# draw Ball in white.
x, y = ball_positions[i]
playing_field[x-BALL_RAD:x+BALL_RAD, y-BALL_RAD:y+BALL_RAD] = white

for (x, y), color in zip([l_paddle_positions[i], r_paddle_positions[i]],
[left_color, right_color]):
# clip y coordinate of the paddle so it does not exceed the screen.
y = max(PADDLE_LEN, y)
y = min(FIELD_IMAGE_SIZE[1] - PADDLE_LEN, y)
y = min(FIELD_SIZE[1] - PADDLE_LEN, y)
playing_field[x:x+PADDLE_WID, y-PADDLE_LEN:y+PADDLE_LEN] = color

# prepare and paste playing field into image
# prepare and paste playing field into the image.
playing_field = np.swapaxes(playing_field, 0, 1)
playing_field = Image.fromarray(playing_field)
background.paste(
playing_field, (int((IMAGE_SIZE[0]-FIELD_IMAGE_SIZE[0])/2), 36))
playing_field, (int((IMAGE_SIZE[0]-FIELD_SIZE[0])/2), 36))

# only draw reward plot and heatmaps every PLOT_INTERVAL iterations or
# every frame when skipping many frames
# only update reward plot and heatmaps every PLOT_INTERVAL iterations
# or every iteration if output_speed is high.
if i % PLOT_INTERVAL == 0 or output_speed > 10:
plt.close()

Expand All @@ -235,7 +240,7 @@ def grayscale_to_heatmap(in_image, min_val, max_val, base_color):
plt.rcParams["figure.autolayout"] = True
plt.rcParams["font.size"] = 8

# set a constant figsize by pixel size
# set figsize by pixel value
DPI = 80
fig = plt.figure(facecolor=background_hex,
figsize=(210/DPI, 115/DPI), dpi=DPI)
Expand All @@ -247,24 +252,25 @@ def grayscale_to_heatmap(in_image, min_val, max_val, base_color):
ax.plot(rewards_right[:i+1], color=right_color/255)
ax.plot(rewards_left[:i+1], color=left_color/255)

# change x_ticks and x_min for the first few plots.
if i < 1600:
x_min = 0
ax.set_xticks(np.arange(0, sim_iterations, 250))
ax.set_xticks(np.arange(0, n_iterations, 250))
else:
x_min = i-1600
ax.set_xticks(np.arange(0, sim_iterations, 500))
ax.set_xticks(np.arange(0, n_iterations, 500))

ax.set_ylabel("mean reward")
# ax.set_xlabel("iteration")
ax.set_yticks([0, 0.5, 1])
ax.set_ylim(0, 1.0)
ax.set_xlim(x_min, i+10)

# fishy workaround to turn a pyplot figure into a PIL.Image
buf = io.BytesIO()
fig.savefig(buf)
buf.seek(0)
reward_plot = Image.open(buf)
# fishy workaround to turn a pyplot figure into a PIL.Image.
buffer = io.BytesIO()
fig.savefig(buffer)
buffer.seek(0)
reward_plot = Image.open(buffer)

background.paste(heatmap_r, (310, 210))
draw.text((315, 290), "weights", tuple(black), font_medium)
Expand All @@ -275,11 +281,11 @@ def grayscale_to_heatmap(in_image, min_val, max_val, base_color):
background.paste(reward_plot, (95, 200))

image_center = int(IMAGE_SIZE[0]/2)
draw.text((image_center - 25, 10), p1_name,
draw.text((image_center - 25, 10), name_left,
tuple(left_color), font_medium, anchor="rt")
draw.text((image_center, 10), "VS", tuple(
black), font_medium, anchor="mt")
draw.text((image_center + 25, 10), p2_name,
draw.text((image_center + 25, 10), name_right,
tuple(right_color), font_medium, anchor="lt")

l_score, r_score = score[i]
Expand All @@ -295,17 +301,17 @@ def grayscale_to_heatmap(in_image, min_val, max_val, base_color):
background.save(os.path.join(temp_dir, f"img_{str(i).zfill(6)}.png"))

# change the speed of the video to show performance before and after
# training at DEFAULT_SPEED and fast forward most of the training
if 75 <= i < 100 or sim_iterations - 300 <= i < sim_iterations - 250:
# training at DEFAULT_SPEED and fast forward most of the training.
if 75 <= i < 100 or n_iterations - 300 <= i < n_iterations - 250:
output_speed = 10
elif 100 <= i < sim_iterations - 300:
elif 100 <= i < n_iterations - 300:
output_speed = 50
else:
output_speed = DEFAULT_SPEED

i += output_speed

print("Images complete, collecting them into a GIF...")
print("Image creation complete, collecting them into a GIF...")

filenames = sorted(glob(os.path.join(temp_dir, "*.png")))

Expand All @@ -320,3 +326,5 @@ def grayscale_to_heatmap(in_image, min_val, max_val, base_color):
for in_file in filenames:
os.unlink(in_file)
os.rmdir(temp_dir)

print("done.")

0 comments on commit 93dfaf2

Please sign in to comment.