forked from hardmaru/write-rnn-tensorflow
-
Notifications
You must be signed in to change notification settings - Fork 1
/
sample_frozen.py
129 lines (95 loc) · 4.35 KB
/
sample_frozen.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
# -*- coding: utf-8 -*-
"""
Created on Thu Feb 23 20:25:16 2017
@author: memo
demonstrates inference with frozen graph def
same as sample.py, but:
- instead of loading model + checkpoint, loads frozen graph
- instead of calling model.sample() function, uses own sample() function with named ops
"""
import numpy as np
import tensorflow as tf
import time
import os
import pickle
import argparse
from utils import *
from model import Model
import random
import svgwrite
from IPython.display import SVG, display
# main code (not in a main function since I want to run this script in IPython as well).
parser = argparse.ArgumentParser()
parser.add_argument('--filename', type=str, default='sample',
help='filename of .svg file to output, without .svg')
parser.add_argument('--sample_length', type=int, default=800,
help='number of strokes to sample')
parser.add_argument('--scale_factor', type=int, default=10,
help='factor to scale down by for svg output. smaller means bigger output')
parser.add_argument('--model_dir', type=str, default='save',
help='directory to save model to')
sample_args = parser.parse_args()
sess = tf.InteractiveSession()
# load frozen graph
from tensorflow.python.platform import gfile
with gfile.FastGFile(os.path.join(sample_args.model_dir, 'graph_frz.pb'),'rb') as f:
graph_def = tf.GraphDef()
graph_def.ParseFromString(f.read())
sess.graph.as_default()
tf.import_graph_def(graph_def, name='')
def sample_stroke():
# don't call model.sample(), instead call sample() function defined below
[strokes, params] = sample(sess, sample_args.sample_length)
draw_strokes(strokes, factor=sample_args.scale_factor, svg_filename = sample_args.filename+'.normal.svg')
draw_strokes_random_color(strokes, factor=sample_args.scale_factor, svg_filename = sample_args.filename+'.color.svg')
draw_strokes_random_color(strokes, factor=sample_args.scale_factor, per_stroke_mode = False, svg_filename = sample_args.filename+'.multi_color.svg')
draw_strokes_eos_weighted(strokes, params, factor=sample_args.scale_factor, svg_filename = sample_args.filename+'.eos_pdf.svg')
draw_strokes_pdf(strokes, params, factor=sample_args.scale_factor, svg_filename = sample_args.filename+'.pdf.svg')
return [strokes, params]
# copied straight from model.sample, but replaced all referenes to 'self' with named ops
def sample(sess, num=1200):
data_in = 'data_in:0'
data_out_pi = 'data_out_pi:0'
data_out_mu1 = 'data_out_mu1:0'
data_out_mu2 = 'data_out_mu2:0'
data_out_sigma1 = 'data_out_sigma1:0'
data_out_sigma2 = 'data_out_sigma2:0'
data_out_corr = 'data_out_corr:0'
data_out_eos = 'data_out_eos:0'
state_in = 'state_in:0'
state_out = 'state_out:0'
def get_pi_idx(x, pdf):
N = pdf.size
accumulate = 0
for i in range(0, N):
accumulate += pdf[i]
if (accumulate >= x):
return i
print('error with sampling ensemble')
return -1
def sample_gaussian_2d(mu1, mu2, s1, s2, rho):
mean = [mu1, mu2]
cov = [[s1*s1, rho*s1*s2], [rho*s1*s2, s2*s2]]
x = np.random.multivariate_normal(mean, cov, 1)
return x[0][0], x[0][1]
prev_x = np.zeros((1, 1, 3), dtype=np.float32)
prev_x[0, 0, 2] = 1 # initially, we want to see beginning of new stroke
prev_state = sess.run(state_in)
strokes = np.zeros((num, 3), dtype=np.float32)
mixture_params = []
for i in range(num):
feed = {data_in: prev_x, state_in:prev_state}
[o_pi, o_mu1, o_mu2, o_sigma1, o_sigma2, o_corr, o_eos, next_state] = sess.run([data_out_pi, data_out_mu1, data_out_mu2, data_out_sigma1, data_out_sigma2, data_out_corr, data_out_eos, state_out],feed)
idx = get_pi_idx(random.random(), o_pi[0])
eos = 1 if random.random() < o_eos[0][0] else 0
next_x1, next_x2 = sample_gaussian_2d(o_mu1[0][idx], o_mu2[0][idx], o_sigma1[0][idx], o_sigma2[0][idx], o_corr[0][idx])
strokes[i,:] = [next_x1, next_x2, eos]
params = [o_pi[0], o_mu1[0], o_mu2[0], o_sigma1[0], o_sigma2[0], o_corr[0], o_eos[0]]
mixture_params.append(params)
prev_x = np.zeros((1, 1, 3), dtype=np.float32)
prev_x[0][0] = np.array([next_x1, next_x2, eos], dtype=np.float32)
prev_state = next_state
strokes[:,0:2] *= 20 #self.args.data_scale # TODO: fix mega hack hardcoding the scale
return strokes, mixture_params
# check output
[strokes, params] = sample_stroke()