mirror of
https://github.com/storytold/artcraft.git
synced 2026-11-02 16:38:14 +00:00
Copy relevant CycleGAN code here.
master @ 1ca8b078c12d915bca1f57e822859fa20bd439c8 Date: Tue Oct 15 10:50:12 2019 -0400
This commit is contained in:
Brandon Thomas committed
1 parent
7c3f3255b0
commit
47022cdeeb
6 files changed
+894
No files matched your search
@@ -0,0 +1,180 @@
|
||||
import os
|
||||
import tensorflow as tf
|
||||
from module import discriminator, generator_gatedcnn
|
||||
from utils import l1_loss, l2_loss, cross_entropy_loss
|
||||
from datetime import datetime
|
||||
|
||||
class CycleGAN(object):
|
||||
|
||||
def __init__(self, num_features, discriminator = discriminator, generator = generator_gatedcnn, mode = 'train', log_dir = './log'):
|
||||
|
||||
self.num_features = num_features
|
||||
self.input_shape = [None, num_features, None] # [batch_size, num_features, num_frames]
|
||||
|
||||
self.discriminator = discriminator
|
||||
self.generator = generator
|
||||
self.mode = mode
|
||||
|
||||
self.build_model()
|
||||
self.optimizer_initializer()
|
||||
|
||||
self.saver = tf.train.Saver()
|
||||
self.sess = tf.Session()
|
||||
self.sess.run(tf.global_variables_initializer())
|
||||
|
||||
if self.mode == 'train':
|
||||
self.train_step = 0
|
||||
now = datetime.now()
|
||||
self.log_dir = os.path.join(log_dir, now.strftime('%Y%m%d-%H%M%S'))
|
||||
self.writer = tf.summary.FileWriter(self.log_dir, tf.get_default_graph())
|
||||
self.generator_summaries, self.discriminator_summaries = self.summary()
|
||||
|
||||
def build_model(self):
|
||||
|
||||
# Placeholders for real training samples
|
||||
self.input_A_real = tf.placeholder(tf.float32, shape = self.input_shape, name = 'input_A_real')
|
||||
self.input_B_real = tf.placeholder(tf.float32, shape = self.input_shape, name = 'input_B_real')
|
||||
# Placeholders for fake generated samples
|
||||
self.input_A_fake = tf.placeholder(tf.float32, shape = self.input_shape, name = 'input_A_fake')
|
||||
self.input_B_fake = tf.placeholder(tf.float32, shape = self.input_shape, name = 'input_B_fake')
|
||||
# Placeholder for test samples
|
||||
self.input_A_test = tf.placeholder(tf.float32, shape = self.input_shape, name = 'input_A_test')
|
||||
self.input_B_test = tf.placeholder(tf.float32, shape = self.input_shape, name = 'input_B_test')
|
||||
|
||||
self.generation_B = self.generator(inputs = self.input_A_real, reuse = False, scope_name = 'generator_A2B')
|
||||
self.cycle_A = self.generator(inputs = self.generation_B, reuse = False, scope_name = 'generator_B2A')
|
||||
|
||||
self.generation_A = self.generator(inputs = self.input_B_real, reuse = True, scope_name = 'generator_B2A')
|
||||
self.cycle_B = self.generator(inputs = self.generation_A, reuse = True, scope_name = 'generator_A2B')
|
||||
|
||||
self.generation_A_identity = self.generator(inputs = self.input_A_real, reuse = True, scope_name = 'generator_B2A')
|
||||
self.generation_B_identity = self.generator(inputs = self.input_B_real, reuse = True, scope_name = 'generator_A2B')
|
||||
|
||||
self.discrimination_A_fake = self.discriminator(inputs = self.generation_A, reuse = False, scope_name = 'discriminator_A')
|
||||
self.discrimination_B_fake = self.discriminator(inputs = self.generation_B, reuse = False, scope_name = 'discriminator_B')
|
||||
|
||||
# Cycle loss
|
||||
self.cycle_loss = l1_loss(y = self.input_A_real, y_hat = self.cycle_A) + l1_loss(y = self.input_B_real, y_hat = self.cycle_B)
|
||||
|
||||
# Identity loss
|
||||
self.identity_loss = l1_loss(y = self.input_A_real, y_hat = self.generation_A_identity) + l1_loss(y = self.input_B_real, y_hat = self.generation_B_identity)
|
||||
|
||||
# Place holder for lambda_cycle and lambda_identity
|
||||
self.lambda_cycle = tf.placeholder(tf.float32, None, name = 'lambda_cycle')
|
||||
self.lambda_identity = tf.placeholder(tf.float32, None, name = 'lambda_identity')
|
||||
|
||||
# Generator loss
|
||||
# Generator wants to fool discriminator
|
||||
self.generator_loss_A2B = l2_loss(y = tf.ones_like(self.discrimination_B_fake), y_hat = self.discrimination_B_fake)
|
||||
self.generator_loss_B2A = l2_loss(y = tf.ones_like(self.discrimination_A_fake), y_hat = self.discrimination_A_fake)
|
||||
|
||||
# Merge the two generators and the cycle loss
|
||||
self.generator_loss = self.generator_loss_A2B + self.generator_loss_B2A + self.lambda_cycle * self.cycle_loss + self.lambda_identity * self.identity_loss
|
||||
|
||||
# Discriminator loss
|
||||
self.discrimination_input_A_real = self.discriminator(inputs = self.input_A_real, reuse = True, scope_name = 'discriminator_A')
|
||||
self.discrimination_input_B_real = self.discriminator(inputs = self.input_B_real, reuse = True, scope_name = 'discriminator_B')
|
||||
self.discrimination_input_A_fake = self.discriminator(inputs = self.input_A_fake, reuse = True, scope_name = 'discriminator_A')
|
||||
self.discrimination_input_B_fake = self.discriminator(inputs = self.input_B_fake, reuse = True, scope_name = 'discriminator_B')
|
||||
|
||||
# Discriminator wants to classify real and fake correctly
|
||||
self.discriminator_loss_input_A_real = l2_loss(y = tf.ones_like(self.discrimination_input_A_real), y_hat = self.discrimination_input_A_real)
|
||||
self.discriminator_loss_input_A_fake = l2_loss(y = tf.zeros_like(self.discrimination_input_A_fake), y_hat = self.discrimination_input_A_fake)
|
||||
self.discriminator_loss_A = (self.discriminator_loss_input_A_real + self.discriminator_loss_input_A_fake) / 2
|
||||
|
||||
self.discriminator_loss_input_B_real = l2_loss(y = tf.ones_like(self.discrimination_input_B_real), y_hat = self.discrimination_input_B_real)
|
||||
self.discriminator_loss_input_B_fake = l2_loss(y = tf.zeros_like(self.discrimination_input_B_fake), y_hat = self.discrimination_input_B_fake)
|
||||
self.discriminator_loss_B = (self.discriminator_loss_input_B_real + self.discriminator_loss_input_B_fake) / 2
|
||||
|
||||
# Merge the two discriminators into one
|
||||
self.discriminator_loss = self.discriminator_loss_A + self.discriminator_loss_B
|
||||
|
||||
# Categorize variables because we have to optimize the two sets of the variables separately
|
||||
trainable_variables = tf.trainable_variables()
|
||||
self.discriminator_vars = [var for var in trainable_variables if 'discriminator' in var.name]
|
||||
self.generator_vars = [var for var in trainable_variables if 'generator' in var.name]
|
||||
#for var in t_vars: print(var.name)
|
||||
|
||||
# Reserved for test
|
||||
self.generation_B_test = self.generator(inputs = self.input_A_test, reuse = True, scope_name = 'generator_A2B')
|
||||
self.generation_A_test = self.generator(inputs = self.input_B_test, reuse = True, scope_name = 'generator_B2A')
|
||||
|
||||
|
||||
def optimizer_initializer(self):
|
||||
|
||||
self.generator_learning_rate = tf.placeholder(tf.float32, None, name = 'generator_learning_rate')
|
||||
self.discriminator_learning_rate = tf.placeholder(tf.float32, None, name = 'discriminator_learning_rate')
|
||||
self.discriminator_optimizer = tf.train.AdamOptimizer(learning_rate = self.discriminator_learning_rate, beta1 = 0.5).minimize(self.discriminator_loss, var_list = self.discriminator_vars)
|
||||
self.generator_optimizer = tf.train.AdamOptimizer(learning_rate = self.generator_learning_rate, beta1 = 0.5).minimize(self.generator_loss, var_list = self.generator_vars)
|
||||
|
||||
def train(self, input_A, input_B, lambda_cycle, lambda_identity, generator_learning_rate, discriminator_learning_rate):
|
||||
|
||||
generation_A, generation_B, generator_loss, _, generator_summaries = self.sess.run(
|
||||
[self.generation_A, self.generation_B, self.generator_loss, self.generator_optimizer, self.generator_summaries], \
|
||||
feed_dict = {self.lambda_cycle: lambda_cycle, self.lambda_identity: lambda_identity, self.input_A_real: input_A, self.input_B_real: input_B, self.generator_learning_rate: generator_learning_rate})
|
||||
|
||||
self.writer.add_summary(generator_summaries, self.train_step)
|
||||
|
||||
discriminator_loss, _, discriminator_summaries = self.sess.run([self.discriminator_loss, self.discriminator_optimizer, self.discriminator_summaries], \
|
||||
feed_dict = {self.input_A_real: input_A, self.input_B_real: input_B, self.discriminator_learning_rate: discriminator_learning_rate, self.input_A_fake: generation_A, self.input_B_fake: generation_B})
|
||||
|
||||
self.writer.add_summary(discriminator_summaries, self.train_step)
|
||||
|
||||
self.train_step += 1
|
||||
|
||||
return generator_loss, discriminator_loss
|
||||
|
||||
|
||||
def test(self, inputs, direction):
|
||||
#print(">>> model.test()!")
|
||||
|
||||
if direction == 'A2B':
|
||||
#print("generation_B_test: {}".format(self.generation_B_test))
|
||||
#print("input_A_test: {}".format(self.input_A_test))
|
||||
generation = self.sess.run(self.generation_B_test, feed_dict = {self.input_A_test: inputs})
|
||||
elif direction == 'B2A':
|
||||
#print("generation_A_test: {}".format(self.generation_A_test))
|
||||
#print("input_B_test: {}".format(self.input_B_test))
|
||||
generation = self.sess.run(self.generation_A_test, feed_dict = {self.input_B_test: inputs})
|
||||
else:
|
||||
raise Exception('Conversion direction must be specified.')
|
||||
|
||||
return generation
|
||||
|
||||
|
||||
def save(self, directory, filename):
|
||||
|
||||
if not os.path.exists(directory):
|
||||
os.makedirs(directory)
|
||||
self.saver.save(self.sess, os.path.join(directory, filename))
|
||||
|
||||
return os.path.join(directory, filename)
|
||||
|
||||
def load(self, filepath):
|
||||
|
||||
self.saver.restore(self.sess, filepath)
|
||||
|
||||
|
||||
def summary(self):
|
||||
|
||||
with tf.name_scope('generator_summaries'):
|
||||
cycle_loss_summary = tf.summary.scalar('cycle_loss', self.cycle_loss)
|
||||
identity_loss_summary = tf.summary.scalar('identity_loss', self.identity_loss)
|
||||
generator_loss_A2B_summary = tf.summary.scalar('generator_loss_A2B', self.generator_loss_A2B)
|
||||
generator_loss_B2A_summary = tf.summary.scalar('generator_loss_B2A', self.generator_loss_B2A)
|
||||
generator_loss_summary = tf.summary.scalar('generator_loss', self.generator_loss)
|
||||
generator_summaries = tf.summary.merge([cycle_loss_summary, identity_loss_summary, generator_loss_A2B_summary, generator_loss_B2A_summary, generator_loss_summary])
|
||||
|
||||
with tf.name_scope('discriminator_summaries'):
|
||||
discriminator_loss_A_summary = tf.summary.scalar('discriminator_loss_A', self.discriminator_loss_A)
|
||||
discriminator_loss_B_summary = tf.summary.scalar('discriminator_loss_B', self.discriminator_loss_B)
|
||||
discriminator_loss_summary = tf.summary.scalar('discriminator_loss', self.discriminator_loss)
|
||||
discriminator_summaries = tf.summary.merge([discriminator_loss_A_summary, discriminator_loss_B_summary, discriminator_loss_summary])
|
||||
|
||||
return generator_summaries, discriminator_summaries
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
|
||||
model = CycleGAN(num_features = 24)
|
||||
#print('Graph Compile Successeded.')
|
||||
@@ -0,0 +1,214 @@
|
||||
import tensorflow as tf
|
||||
|
||||
def gated_linear_layer(inputs, gates, name = None):
|
||||
|
||||
activation = tf.multiply(x = inputs, y = tf.sigmoid(gates), name = name)
|
||||
|
||||
return activation
|
||||
|
||||
def instance_norm_layer(
|
||||
inputs,
|
||||
epsilon = 1e-06,
|
||||
activation_fn = None,
|
||||
name = None):
|
||||
|
||||
instance_norm_layer = tf.contrib.layers.instance_norm(
|
||||
inputs = inputs,
|
||||
epsilon = epsilon,
|
||||
activation_fn = activation_fn)
|
||||
|
||||
return instance_norm_layer
|
||||
|
||||
def conv1d_layer(
|
||||
inputs,
|
||||
filters,
|
||||
kernel_size,
|
||||
strides = 1,
|
||||
padding = 'same',
|
||||
activation = None,
|
||||
kernel_initializer = None,
|
||||
name = None):
|
||||
|
||||
conv_layer = tf.layers.conv1d(
|
||||
inputs = inputs,
|
||||
filters = filters,
|
||||
kernel_size = kernel_size,
|
||||
strides = strides,
|
||||
padding = padding,
|
||||
activation = activation,
|
||||
kernel_initializer = kernel_initializer,
|
||||
name = name)
|
||||
|
||||
return conv_layer
|
||||
|
||||
def conv2d_layer(
|
||||
inputs,
|
||||
filters,
|
||||
kernel_size,
|
||||
strides,
|
||||
padding = 'same',
|
||||
activation = None,
|
||||
kernel_initializer = None,
|
||||
name = None):
|
||||
|
||||
conv_layer = tf.layers.conv2d(
|
||||
inputs = inputs,
|
||||
filters = filters,
|
||||
kernel_size = kernel_size,
|
||||
strides = strides,
|
||||
padding = padding,
|
||||
activation = activation,
|
||||
kernel_initializer = kernel_initializer,
|
||||
name = name)
|
||||
|
||||
return conv_layer
|
||||
|
||||
def residual1d_block(
|
||||
inputs,
|
||||
filters = 1024,
|
||||
kernel_size = 3,
|
||||
strides = 1,
|
||||
name_prefix = 'residule_block_'):
|
||||
|
||||
h1 = conv1d_layer(inputs = inputs, filters = filters, kernel_size = kernel_size, strides = strides, activation = None, name = name_prefix + 'h1_conv')
|
||||
h1_norm = instance_norm_layer(inputs = h1, activation_fn = None, name = name_prefix + 'h1_norm')
|
||||
h1_gates = conv1d_layer(inputs = inputs, filters = filters, kernel_size = kernel_size, strides = strides, activation = None, name = name_prefix + 'h1_gates')
|
||||
h1_norm_gates = instance_norm_layer(inputs = h1_gates, activation_fn = None, name = name_prefix + 'h1_norm_gates')
|
||||
h1_glu = gated_linear_layer(inputs = h1_norm, gates = h1_norm_gates, name = name_prefix + 'h1_glu')
|
||||
h2 = conv1d_layer(inputs = h1_glu, filters = filters // 2, kernel_size = kernel_size, strides = strides, activation = None, name = name_prefix + 'h2_conv')
|
||||
h2_norm = instance_norm_layer(inputs = h2, activation_fn = None, name = name_prefix + 'h2_norm')
|
||||
|
||||
h3 = inputs + h2_norm
|
||||
|
||||
return h3
|
||||
|
||||
def downsample1d_block(
|
||||
inputs,
|
||||
filters,
|
||||
kernel_size,
|
||||
strides,
|
||||
name_prefix = 'downsample1d_block_'):
|
||||
|
||||
h1 = conv1d_layer(inputs = inputs, filters = filters, kernel_size = kernel_size, strides = strides, activation = None, name = name_prefix + 'h1_conv')
|
||||
h1_norm = instance_norm_layer(inputs = h1, activation_fn = None, name = name_prefix + 'h1_norm')
|
||||
h1_gates = conv1d_layer(inputs = inputs, filters = filters, kernel_size = kernel_size, strides = strides, activation = None, name = name_prefix + 'h1_gates')
|
||||
h1_norm_gates = instance_norm_layer(inputs = h1_gates, activation_fn = None, name = name_prefix + 'h1_norm_gates')
|
||||
h1_glu = gated_linear_layer(inputs = h1_norm, gates = h1_norm_gates, name = name_prefix + 'h1_glu')
|
||||
|
||||
return h1_glu
|
||||
|
||||
def downsample2d_block(
|
||||
inputs,
|
||||
filters,
|
||||
kernel_size,
|
||||
strides,
|
||||
name_prefix = 'downsample2d_block_'):
|
||||
|
||||
h1 = conv2d_layer(inputs = inputs, filters = filters, kernel_size = kernel_size, strides = strides, activation = None, name = name_prefix + 'h1_conv')
|
||||
h1_norm = instance_norm_layer(inputs = h1, activation_fn = None, name = name_prefix + 'h1_norm')
|
||||
h1_gates = conv2d_layer(inputs = inputs, filters = filters, kernel_size = kernel_size, strides = strides, activation = None, name = name_prefix + 'h1_gates')
|
||||
h1_norm_gates = instance_norm_layer(inputs = h1_gates, activation_fn = None, name = name_prefix + 'h1_norm_gates')
|
||||
h1_glu = gated_linear_layer(inputs = h1_norm, gates = h1_norm_gates, name = name_prefix + 'h1_glu')
|
||||
|
||||
return h1_glu
|
||||
|
||||
def upsample1d_block(
|
||||
inputs,
|
||||
filters,
|
||||
kernel_size,
|
||||
strides,
|
||||
shuffle_size = 2,
|
||||
name_prefix = 'upsample1d_block_'):
|
||||
|
||||
h1 = conv1d_layer(inputs = inputs, filters = filters, kernel_size = kernel_size, strides = strides, activation = None, name = name_prefix + 'h1_conv')
|
||||
h1_shuffle = pixel_shuffler(inputs = h1, shuffle_size = shuffle_size, name = name_prefix + 'h1_shuffle')
|
||||
h1_norm = instance_norm_layer(inputs = h1_shuffle, activation_fn = None, name = name_prefix + 'h1_norm')
|
||||
|
||||
h1_gates = conv1d_layer(inputs = inputs, filters = filters, kernel_size = kernel_size, strides = strides, activation = None, name = name_prefix + 'h1_gates')
|
||||
h1_shuffle_gates = pixel_shuffler(inputs = h1_gates, shuffle_size = shuffle_size, name = name_prefix + 'h1_shuffle_gates')
|
||||
h1_norm_gates = instance_norm_layer(inputs = h1_shuffle_gates, activation_fn = None, name = name_prefix + 'h1_norm_gates')
|
||||
|
||||
h1_glu = gated_linear_layer(inputs = h1_norm, gates = h1_norm_gates, name = name_prefix + 'h1_glu')
|
||||
|
||||
return h1_glu
|
||||
|
||||
def pixel_shuffler(inputs, shuffle_size = 2, name = None):
|
||||
|
||||
n = tf.shape(inputs)[0]
|
||||
w = tf.shape(inputs)[1]
|
||||
c = inputs.get_shape().as_list()[2]
|
||||
|
||||
oc = c // shuffle_size
|
||||
ow = w * shuffle_size
|
||||
|
||||
outputs = tf.reshape(tensor = inputs, shape = [n, ow, oc], name = name)
|
||||
|
||||
return outputs
|
||||
|
||||
def generator_gatedcnn(inputs, reuse = False, scope_name = 'generator_gatedcnn'):
|
||||
|
||||
# inputs has shape [batch_size, num_features, time]
|
||||
# we need to convert it to [batch_size, time, num_features] for 1D convolution
|
||||
inputs = tf.transpose(inputs, perm = [0, 2, 1], name = 'input_transpose')
|
||||
|
||||
with tf.variable_scope(scope_name) as scope:
|
||||
# Discriminator would be reused in CycleGAN
|
||||
if reuse:
|
||||
scope.reuse_variables()
|
||||
else:
|
||||
assert scope.reuse is False
|
||||
|
||||
h1 = conv1d_layer(inputs = inputs, filters = 128, kernel_size = 15, strides = 1, activation = None, name = 'h1_conv')
|
||||
h1_gates = conv1d_layer(inputs = inputs, filters = 128, kernel_size = 15, strides = 1, activation = None, name = 'h1_conv_gates')
|
||||
h1_glu = gated_linear_layer(inputs = h1, gates = h1_gates, name = 'h1_glu')
|
||||
|
||||
# Downsample
|
||||
d1 = downsample1d_block(inputs = h1_glu, filters = 256, kernel_size = 5, strides = 2, name_prefix = 'downsample1d_block1_')
|
||||
d2 = downsample1d_block(inputs = d1, filters = 512, kernel_size = 5, strides = 2, name_prefix = 'downsample1d_block2_')
|
||||
|
||||
# Residual blocks
|
||||
r1 = residual1d_block(inputs = d2, filters = 1024, kernel_size = 3, strides = 1, name_prefix = 'residual1d_block1_')
|
||||
r2 = residual1d_block(inputs = r1, filters = 1024, kernel_size = 3, strides = 1, name_prefix = 'residual1d_block2_')
|
||||
r3 = residual1d_block(inputs = r2, filters = 1024, kernel_size = 3, strides = 1, name_prefix = 'residual1d_block3_')
|
||||
r4 = residual1d_block(inputs = r3, filters = 1024, kernel_size = 3, strides = 1, name_prefix = 'residual1d_block4_')
|
||||
r5 = residual1d_block(inputs = r4, filters = 1024, kernel_size = 3, strides = 1, name_prefix = 'residual1d_block5_')
|
||||
r6 = residual1d_block(inputs = r5, filters = 1024, kernel_size = 3, strides = 1, name_prefix = 'residual1d_block6_')
|
||||
|
||||
# Upsample
|
||||
u1 = upsample1d_block(inputs = r6, filters = 1024, kernel_size = 5, strides = 1, shuffle_size = 2, name_prefix = 'upsample1d_block1_')
|
||||
u2 = upsample1d_block(inputs = u1, filters = 512, kernel_size = 5, strides = 1, shuffle_size = 2, name_prefix = 'upsample1d_block2_')
|
||||
|
||||
# Output
|
||||
o1 = conv1d_layer(inputs = u2, filters = 24, kernel_size = 15, strides = 1, activation = None, name = 'o1_conv')
|
||||
o2 = tf.transpose(o1, perm = [0, 2, 1], name = 'output_transpose')
|
||||
|
||||
return o2
|
||||
|
||||
|
||||
def discriminator(inputs, reuse = False, scope_name = 'discriminator'):
|
||||
|
||||
# inputs has shape [batch_size, num_features, time]
|
||||
# we need to add channel for 2D convolution [batch_size, num_features, time, 1]
|
||||
inputs = tf.expand_dims(inputs, -1)
|
||||
|
||||
with tf.variable_scope(scope_name) as scope:
|
||||
# Discriminator would be reused in CycleGAN
|
||||
if reuse:
|
||||
scope.reuse_variables()
|
||||
else:
|
||||
assert scope.reuse is False
|
||||
|
||||
h1 = conv2d_layer(inputs = inputs, filters = 128, kernel_size = [3, 3], strides = [1, 2], activation = None, name = 'h1_conv')
|
||||
h1_gates = conv2d_layer(inputs = inputs, filters = 128, kernel_size = [3, 3], strides = [1, 2], activation = None, name = 'h1_conv_gates')
|
||||
h1_glu = gated_linear_layer(inputs = h1, gates = h1_gates, name = 'h1_glu')
|
||||
|
||||
# Downsample
|
||||
d1 = downsample2d_block(inputs = h1_glu, filters = 256, kernel_size = [3, 3], strides = [2, 2], name_prefix = 'downsample2d_block1_')
|
||||
d2 = downsample2d_block(inputs = d1, filters = 512, kernel_size = [3, 3], strides = [2, 2], name_prefix = 'downsample2d_block2_')
|
||||
d3 = downsample2d_block(inputs = d2, filters = 1024, kernel_size = [6, 3], strides = [1, 2], name_prefix = 'downsample2d_block3_')
|
||||
|
||||
# Output
|
||||
o1 = tf.layers.dense(inputs = d3, units = 1, activation = tf.nn.sigmoid)
|
||||
|
||||
return o1
|
||||
|
||||
@@ -0,0 +1,238 @@
|
||||
import librosa
|
||||
import numpy as np
|
||||
import os
|
||||
import pyworld
|
||||
|
||||
def load_wavs(wav_dir, sr):
|
||||
|
||||
wavs = list()
|
||||
for file in os.listdir(wav_dir):
|
||||
file_path = os.path.join(wav_dir, file)
|
||||
wav, _ = librosa.load(file_path, sr = sr, mono = True)
|
||||
#wav = wav.astype(np.float64)
|
||||
wavs.append(wav)
|
||||
|
||||
return wavs
|
||||
|
||||
def world_decompose(wav, fs, frame_period = 5.0):
|
||||
|
||||
# Decompose speech signal into f0, spectral envelope and aperiodicity using WORLD
|
||||
wav = wav.astype(np.float64)
|
||||
f0, timeaxis = pyworld.harvest(wav, fs, frame_period = frame_period, f0_floor = 71.0, f0_ceil = 800.0)
|
||||
sp = pyworld.cheaptrick(wav, f0, timeaxis, fs)
|
||||
ap = pyworld.d4c(wav, f0, timeaxis, fs)
|
||||
|
||||
return f0, timeaxis, sp, ap
|
||||
|
||||
def world_encode_spectral_envelop(sp, fs, dim = 24):
|
||||
|
||||
# Get Mel-cepstral coefficients (MCEPs)
|
||||
|
||||
#sp = sp.astype(np.float64)
|
||||
coded_sp = pyworld.code_spectral_envelope(sp, fs, dim)
|
||||
|
||||
return coded_sp
|
||||
|
||||
def world_decode_spectral_envelop(coded_sp, fs):
|
||||
|
||||
fftlen = pyworld.get_cheaptrick_fft_size(fs)
|
||||
#coded_sp = coded_sp.astype(np.float32)
|
||||
#coded_sp = np.ascontiguousarray(coded_sp)
|
||||
decoded_sp = pyworld.decode_spectral_envelope(coded_sp, fs, fftlen)
|
||||
|
||||
return decoded_sp
|
||||
|
||||
|
||||
def world_encode_data(wavs, fs, frame_period = 5.0, coded_dim = 24):
|
||||
|
||||
f0s = list()
|
||||
timeaxes = list()
|
||||
sps = list()
|
||||
aps = list()
|
||||
coded_sps = list()
|
||||
|
||||
for wav in wavs:
|
||||
f0, timeaxis, sp, ap = world_decompose(wav = wav, fs = fs, frame_period = frame_period)
|
||||
coded_sp = world_encode_spectral_envelop(sp = sp, fs = fs, dim = coded_dim)
|
||||
f0s.append(f0)
|
||||
timeaxes.append(timeaxis)
|
||||
sps.append(sp)
|
||||
aps.append(ap)
|
||||
coded_sps.append(coded_sp)
|
||||
|
||||
return f0s, timeaxes, sps, aps, coded_sps
|
||||
|
||||
|
||||
def transpose_in_list(lst):
|
||||
|
||||
transposed_lst = list()
|
||||
for array in lst:
|
||||
transposed_lst.append(array.T)
|
||||
return transposed_lst
|
||||
|
||||
|
||||
def world_decode_data(coded_sps, fs):
|
||||
|
||||
decoded_sps = list()
|
||||
|
||||
for coded_sp in coded_sps:
|
||||
decoded_sp = world_decode_spectral_envelop(coded_sp, fs)
|
||||
decoded_sps.append(decoded_sp)
|
||||
|
||||
return decoded_sps
|
||||
|
||||
|
||||
def world_speech_synthesis(f0, decoded_sp, ap, fs, frame_period):
|
||||
|
||||
#decoded_sp = decoded_sp.astype(np.float64)
|
||||
wav = pyworld.synthesize(f0, decoded_sp, ap, fs, frame_period)
|
||||
# Librosa could not save wav if not doing so
|
||||
wav = wav.astype(np.float32)
|
||||
|
||||
return wav
|
||||
|
||||
|
||||
def world_synthesis_data(f0s, decoded_sps, aps, fs, frame_period):
|
||||
|
||||
wavs = list()
|
||||
|
||||
for f0, decoded_sp, ap in zip(f0s, decoded_sps, aps):
|
||||
wav = world_speech_synthesis(f0, decoded_sp, ap, fs, frame_period)
|
||||
wavs.append(wav)
|
||||
|
||||
return wavs
|
||||
|
||||
|
||||
def coded_sps_normalization_fit_transoform(coded_sps):
|
||||
|
||||
coded_sps_concatenated = np.concatenate(coded_sps, axis = 1)
|
||||
coded_sps_mean = np.mean(coded_sps_concatenated, axis = 1, keepdims = True)
|
||||
coded_sps_std = np.std(coded_sps_concatenated, axis = 1, keepdims = True)
|
||||
|
||||
coded_sps_normalized = list()
|
||||
for coded_sp in coded_sps:
|
||||
coded_sps_normalized.append((coded_sp - coded_sps_mean) / coded_sps_std)
|
||||
|
||||
return coded_sps_normalized, coded_sps_mean, coded_sps_std
|
||||
|
||||
def coded_sps_normalization_transoform(coded_sps, coded_sps_mean, coded_sps_std):
|
||||
|
||||
coded_sps_normalized = list()
|
||||
for coded_sp in coded_sps:
|
||||
coded_sps_normalized.append((coded_sp - coded_sps_mean) / coded_sps_std)
|
||||
|
||||
return coded_sps_normalized
|
||||
|
||||
def coded_sps_normalization_inverse_transoform(normalized_coded_sps, coded_sps_mean, coded_sps_std):
|
||||
|
||||
coded_sps = list()
|
||||
for normalized_coded_sp in normalized_coded_sps:
|
||||
coded_sps.append(normalized_coded_sp * coded_sps_std + coded_sps_mean)
|
||||
|
||||
return coded_sps
|
||||
|
||||
def coded_sp_padding(coded_sp, multiple = 4):
|
||||
|
||||
num_features = coded_sp.shape[0]
|
||||
num_frames = coded_sp.shape[1]
|
||||
num_frames_padded = int(np.ceil(num_frames / multiple)) * multiple
|
||||
num_frames_diff = num_frames_padded - num_frames
|
||||
num_pad_left = num_frames_diff // 2
|
||||
num_pad_right = num_frames_diff - num_pad_left
|
||||
coded_sp_padded = np.pad(coded_sp, ((0, 0), (num_pad_left, num_pad_right)), 'constant', constant_values = 0)
|
||||
|
||||
return coded_sp_padded
|
||||
|
||||
def wav_padding(wav, sr, frame_period, multiple = 4):
|
||||
|
||||
assert wav.ndim == 1
|
||||
num_frames = len(wav)
|
||||
num_frames_padded = int((np.ceil((np.floor(num_frames / (sr * frame_period / 1000)) + 1) / multiple + 1) * multiple - 1) * (sr * frame_period / 1000))
|
||||
num_frames_diff = num_frames_padded - num_frames
|
||||
num_pad_left = num_frames_diff // 2
|
||||
num_pad_right = num_frames_diff - num_pad_left
|
||||
wav_padded = np.pad(wav, (num_pad_left, num_pad_right), 'constant', constant_values = 0)
|
||||
|
||||
return wav_padded
|
||||
|
||||
|
||||
def logf0_statistics(f0s):
|
||||
|
||||
log_f0s_concatenated = np.ma.log(np.concatenate(f0s))
|
||||
log_f0s_mean = log_f0s_concatenated.mean()
|
||||
log_f0s_std = log_f0s_concatenated.std()
|
||||
|
||||
return log_f0s_mean, log_f0s_std
|
||||
|
||||
def pitch_conversion(f0, mean_log_src, std_log_src, mean_log_target, std_log_target):
|
||||
|
||||
# Logarithm Gaussian normalization for Pitch Conversions
|
||||
f0_converted = np.exp((np.log(f0) - mean_log_src) / std_log_src * std_log_target + mean_log_target)
|
||||
|
||||
return f0_converted
|
||||
|
||||
def wavs_to_specs(wavs, n_fft = 1024, hop_length = None):
|
||||
|
||||
stfts = list()
|
||||
for wav in wavs:
|
||||
stft = librosa.stft(wav, n_fft = n_fft, hop_length = hop_length)
|
||||
stfts.append(stft)
|
||||
|
||||
return stfts
|
||||
|
||||
|
||||
def wavs_to_mfccs(wavs, sr, n_fft = 1024, hop_length = None, n_mels = 128, n_mfcc = 24):
|
||||
|
||||
mfccs = list()
|
||||
for wav in wavs:
|
||||
mfcc = librosa.feature.mfcc(y = wav, sr = sr, n_fft = n_fft, hop_length = hop_length, n_mels = n_mels, n_mfcc = n_mfcc)
|
||||
mfccs.append(mfcc)
|
||||
|
||||
return mfccs
|
||||
|
||||
|
||||
def mfccs_normalization(mfccs):
|
||||
|
||||
mfccs_concatenated = np.concatenate(mfccs, axis = 1)
|
||||
mfccs_mean = np.mean(mfccs_concatenated, axis = 1, keepdims = True)
|
||||
mfccs_std = np.std(mfccs_concatenated, axis = 1, keepdims = True)
|
||||
|
||||
mfccs_normalized = list()
|
||||
for mfcc in mfccs:
|
||||
mfccs_normalized.append((mfcc - mfccs_mean) / mfccs_std)
|
||||
|
||||
return mfccs_normalized, mfccs_mean, mfccs_std
|
||||
|
||||
|
||||
def sample_train_data(dataset_A, dataset_B, n_frames = 128):
|
||||
|
||||
num_samples = min(len(dataset_A), len(dataset_B))
|
||||
train_data_A_idx = np.arange(len(dataset_A))
|
||||
train_data_B_idx = np.arange(len(dataset_B))
|
||||
np.random.shuffle(train_data_A_idx)
|
||||
np.random.shuffle(train_data_B_idx)
|
||||
train_data_A_idx_subset = train_data_A_idx[:num_samples]
|
||||
train_data_B_idx_subset = train_data_B_idx[:num_samples]
|
||||
|
||||
train_data_A = list()
|
||||
train_data_B = list()
|
||||
|
||||
for idx_A, idx_B in zip(train_data_A_idx_subset, train_data_B_idx_subset):
|
||||
data_A = dataset_A[idx_A]
|
||||
frames_A_total = data_A.shape[1]
|
||||
assert frames_A_total >= n_frames
|
||||
start_A = np.random.randint(frames_A_total - n_frames + 1)
|
||||
end_A = start_A + n_frames
|
||||
train_data_A.append(data_A[:,start_A:end_A])
|
||||
|
||||
data_B = dataset_B[idx_B]
|
||||
frames_B_total = data_B.shape[1]
|
||||
assert frames_B_total >= n_frames
|
||||
start_B = np.random.randint(frames_B_total - n_frames + 1)
|
||||
end_B = start_B + n_frames
|
||||
train_data_B.append(data_B[:,start_B:end_B])
|
||||
|
||||
train_data_A = np.array(train_data_A)
|
||||
train_data_B = np.array(train_data_B)
|
||||
|
||||
return train_data_A, train_data_B
|
||||
@@ -0,0 +1,36 @@
|
||||
absl-py==0.7.1
|
||||
astor==0.8.0
|
||||
audioread==2.1.8
|
||||
cffi==1.12.3
|
||||
Cython==0.29.12
|
||||
decorator==4.4.0
|
||||
falcon==2.0.0
|
||||
falcon-multipart==0.2.0
|
||||
gast==0.2.2
|
||||
google-pasta==0.1.7
|
||||
grpcio==1.22.0
|
||||
h5py==2.9.0
|
||||
joblib==0.13.2
|
||||
Keras-Applications==1.0.8
|
||||
Keras-Preprocessing==1.1.0
|
||||
librosa==0.7.0
|
||||
llvmlite==0.29.0
|
||||
Markdown==3.1.1
|
||||
numba==0.45.0
|
||||
numpy==1.16.4
|
||||
protobuf==3.9.0
|
||||
pycparser==2.19
|
||||
pyworld==0.2.8
|
||||
pyzmq==18.1.0
|
||||
resampy==0.2.1
|
||||
scikit-learn==0.21.2
|
||||
scipy==1.3.0
|
||||
six==1.12.0
|
||||
SoundFile==0.10.2
|
||||
tensorboard==1.14.0
|
||||
tensorflow-estimator==1.14.0
|
||||
tensorflow-gpu==1.14.0
|
||||
termcolor==1.1.0
|
||||
Werkzeug==0.15.5
|
||||
wrapt==1.11.2
|
||||
zmq==0.0.0
|
||||
Executable
+209
@@ -0,0 +1,209 @@
|
||||
#!/usr/bin/env python3
|
||||
|
||||
"""
|
||||
Starts a demo HTTP server to capture and transform audio
|
||||
as a live demonstration of the trained model.
|
||||
|
||||
Brandon Thomas 2019-07-29 <bt@brand.io> <echelon@gmail.com>
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import io
|
||||
import librosa
|
||||
import numpy as np
|
||||
import os
|
||||
import pathlib
|
||||
import scipy
|
||||
import soundfile
|
||||
import struct
|
||||
import subprocess
|
||||
import tempfile
|
||||
import tensorflow as tf
|
||||
import zmq
|
||||
|
||||
from falcon_multipart.middleware import MultipartMiddleware
|
||||
from model import CycleGAN
|
||||
from preprocess import *
|
||||
from wsgiref import simple_server
|
||||
|
||||
print("TensorFlow version: {}".format(tf.version.VERSION))
|
||||
|
||||
|
||||
class Converter():
|
||||
def __init__(self, model_dir, model_name):
|
||||
self.num_features = 24
|
||||
self.sampling_rate = 16000
|
||||
self.frame_period = 5.0
|
||||
|
||||
self.model = CycleGAN(num_features = self.num_features, mode = 'test')
|
||||
|
||||
self.model.load(filepath = os.path.join(model_dir, model_name))
|
||||
|
||||
self.mcep_normalization_params = np.load(os.path.join(model_dir, 'mcep_normalization.npz'))
|
||||
self.mcep_mean_A = self.mcep_normalization_params['mean_A']
|
||||
self.mcep_std_A = self.mcep_normalization_params['std_A']
|
||||
self.mcep_mean_B = self.mcep_normalization_params['mean_B']
|
||||
self.mcep_std_B = self.mcep_normalization_params['std_B']
|
||||
|
||||
self.logf0s_normalization_params = np.load(os.path.join(model_dir,
|
||||
'logf0s_normalization.npz'))
|
||||
self.logf0s_mean_A = self.logf0s_normalization_params['mean_A']
|
||||
self.logf0s_std_A = self.logf0s_normalization_params['std_A']
|
||||
self.logf0s_mean_B = self.logf0s_normalization_params['mean_B']
|
||||
self.logf0s_std_B = self.logf0s_normalization_params['std_B']
|
||||
|
||||
def convert_partial(self, wav, conversion_direction='A2B'):
|
||||
wav = wav_padding(wav = wav,
|
||||
sr = self.sampling_rate,
|
||||
frame_period = self.frame_period,
|
||||
multiple = 4)
|
||||
f0, timeaxis, sp, ap = world_decompose(wav = wav,
|
||||
fs = self.sampling_rate,
|
||||
frame_period = self.frame_period)
|
||||
coded_sp = world_encode_spectral_envelop(sp = sp,
|
||||
fs = self.sampling_rate,
|
||||
dim = self.num_features)
|
||||
coded_sp_transposed = coded_sp.T
|
||||
|
||||
f0_converted = pitch_conversion(f0 = f0,
|
||||
mean_log_src = self.logf0s_mean_A,
|
||||
std_log_src = self.logf0s_std_A,
|
||||
mean_log_target = self.logf0s_mean_B,
|
||||
std_log_target = self.logf0s_std_B)
|
||||
coded_sp_norm = (coded_sp_transposed - self.mcep_mean_A) / self.mcep_std_A
|
||||
|
||||
coded_sp_converted_norm = self.model.test(inputs = np.array([coded_sp_norm]),
|
||||
direction = conversion_direction)[0]
|
||||
coded_sp_converted = coded_sp_converted_norm * self.mcep_std_B + self.mcep_mean_B
|
||||
|
||||
coded_sp_converted = coded_sp_converted.T
|
||||
coded_sp_converted = np.ascontiguousarray(coded_sp_converted)
|
||||
decoded_sp_converted = world_decode_spectral_envelop(coded_sp = coded_sp_converted,
|
||||
fs = self.sampling_rate)
|
||||
wav_transformed = world_speech_synthesis(f0 = f0_converted,
|
||||
decoded_sp = decoded_sp_converted,
|
||||
ap = ap,
|
||||
fs = self.sampling_rate,
|
||||
frame_period = self.frame_period)
|
||||
|
||||
# For debugging model output, uncomment the following line:
|
||||
# librosa.output.write_wav('model_output.wav', wav_transformed, self.sampling_rate)
|
||||
|
||||
"""
|
||||
# TODO: Perhaps ditch this. It's probably unnecessary work.
|
||||
upsampled = librosa.resample(wav_transformed, self.sampling_rate, 48000)
|
||||
pcm_data = upsampled.astype(np.float64)
|
||||
stereo_pcm_data = np.tile(pcm_data, (2,1)).T
|
||||
return stereo_pcm_data.astype(np.float32)
|
||||
"""
|
||||
|
||||
#return wav
|
||||
return wav_transformed
|
||||
|
||||
#def convert(self, wav, conversion_direction='A2B'):
|
||||
# pcm_data = self.convert_partial(wav, conversion_direction=conversion_direction)
|
||||
# buf = io.BytesIO()
|
||||
# # pcm_data: A 1-D or 2-D numpy array of either integer or float data-type.
|
||||
# # To write multiple-channels, use a 2-D array of shape (Nsamples, Nchannels).
|
||||
# scipy.io.wavfile.write(buf, 48000, pcm_data)
|
||||
# return buf
|
||||
|
||||
# Set up model
|
||||
# This should live long in memory, so we do it up front.
|
||||
model_dir_default = './model/sf1_tm1'
|
||||
model_name_default = 'sf1_tm1.ckpt'
|
||||
|
||||
# TODO: UNCOMMENT
|
||||
converter = Converter(model_dir_default, model_name_default)
|
||||
|
||||
TEMP_DIR = tempfile.TemporaryDirectory(prefix='queue_audio')
|
||||
|
||||
def temp_file_name(suffix='.wav'):
|
||||
# NB: Not actually using the tempfile. Just the random name.
|
||||
temp_file = tempfile.NamedTemporaryFile(suffix=suffix)
|
||||
name = os.path.basename(temp_file.name)
|
||||
return os.path.join(TEMP_DIR.name, name)
|
||||
|
||||
SAVE_FILES = False
|
||||
SKIP_CONVERT = False
|
||||
|
||||
def convert(audio):
|
||||
#audio = np.array(audio, dtype=np.int16)
|
||||
audio = np.array(audio, dtype=np.float32)
|
||||
print(audio)
|
||||
|
||||
#source_rate = 38000 # Experimentally determined for Rust library 'wavy'
|
||||
#source_rate = 88000 # Experimentally determined for Rust library 'CPAL'
|
||||
source_rate = 44100
|
||||
|
||||
if SAVE_FILES:
|
||||
filename = temp_file_name('.wav')
|
||||
print('----- Original wav file out: {}'.format(filename))
|
||||
scipy.io.wavfile.write(filename, source_rate, audio)
|
||||
|
||||
audio = librosa.resample(audio, source_rate, 16000)
|
||||
|
||||
if SAVE_FILES:
|
||||
filename = temp_file_name('.wav')
|
||||
print('----- Downsampled file out: {}'.format(filename))
|
||||
scipy.io.wavfile.write(filename, 16000, audio)
|
||||
|
||||
if SKIP_CONVERT:
|
||||
return
|
||||
|
||||
results = converter.convert_partial(audio, conversion_direction = 'A2B')
|
||||
#results = audio[:]
|
||||
|
||||
#consume_rate = 68000 # Experimentally determined for Rust lib 'CPAL'
|
||||
consume_rate = 44100
|
||||
upsampled = librosa.resample(results, 16000, consume_rate)
|
||||
|
||||
if SAVE_FILES:
|
||||
filename = temp_file_name('.wav')
|
||||
print('----- Upsampled (transformed) file out: {}'.format(filename))
|
||||
scipy.io.wavfile.write(filename, consume_rate, upsampled)
|
||||
|
||||
return upsampled
|
||||
|
||||
|
||||
BUFFER_SIZE = 5000
|
||||
MULTIPLIER = 10
|
||||
#MULTIPLIER = 1
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument('--port', type=int, default=5555)
|
||||
args = parser.parse_args()
|
||||
|
||||
context = zmq.Context()
|
||||
socket = context.socket(zmq.REP)
|
||||
socket.bind("tcp://*:{}".format(args.port))
|
||||
print('Running server...')
|
||||
|
||||
queue = []
|
||||
while True:
|
||||
# Wait for next request from client
|
||||
message = socket.recv()
|
||||
#socket.send(message)
|
||||
#continue
|
||||
|
||||
# 16-bit PCM (-32768, +32767) int16
|
||||
# layout = '!' + ('h' * BUFFER_SIZE) # NB: Audio from 'wavy' Rust library
|
||||
layout = '<' + ('f' * BUFFER_SIZE) # NB: Audio from 'CPAL' Rust library
|
||||
decoded = struct.unpack(layout, message)
|
||||
queue.extend(decoded)
|
||||
|
||||
if len(queue) >= BUFFER_SIZE * MULTIPLIER:
|
||||
#results = queue[:]
|
||||
results = convert(queue)
|
||||
queue = []
|
||||
layout = '<' + ('f' * len(results))
|
||||
result_bytes = struct.pack(layout, *results)
|
||||
socket.send(result_bytes)
|
||||
else:
|
||||
# Must send reply back to client
|
||||
socket.send(b"OK")
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
|
||||
@@ -0,0 +1,17 @@
|
||||
import tensorflow as tf
|
||||
import os
|
||||
import random
|
||||
import numpy as np
|
||||
|
||||
def l1_loss(y, y_hat):
|
||||
|
||||
return tf.reduce_mean(tf.abs(y - y_hat))
|
||||
|
||||
def l2_loss(y, y_hat):
|
||||
|
||||
return tf.reduce_mean(tf.square(y - y_hat))
|
||||
|
||||
def cross_entropy_loss(logits, labels):
|
||||
return tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(logits = logits, labels = labels))
|
||||
|
||||
|
||||
Reference in new issue
Block a user