visualize_vgg_model.py 3.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107
  1. import numpy as np
  2. from scipy import misc
  3. import tensorflow as tf
  4. from matplotlib import pyplot, image
  5. import vggverydeep19
  6. paintingStyleImage = image.imread("../data/schoolofathens.jpg")
  7. pyplot.imshow(paintingStyleImage)
  8. inputImage = image.imread("../data/grandcentral.jpg")
  9. pyplot.imshow(inputImage)
  10. outputWidth = 800
  11. outputHeight = 600
  12. # Beta constant
  13. beta = 5
  14. # Alpha constant
  15. alpha = 100
  16. # Noise ratio
  17. noiseRatio = 0.6
  18. nodes = vggverydeep19.load('../data/imagenet-vgg-verydeep-19.mat', (600, 800))
  19. # Mean VGG-19 image
  20. meanImage19 = np.array([103.939, 116.779, 123.68]).reshape((1,1,1,3)) #pylint: disable=no-member
  21. # Squared-error loss of content between the two feature representations
  22. def sqErrorLossContent(sess, modelGraph, layer):
  23. p = session.run(modelGraph[layer])
  24. #pylint: disable=maybe-no-member
  25. N = p.shape[3]
  26. M = p.shape[1] * p.shape[2]
  27. return (1 / (4 * N * M)) * tf.reduce_sum(tf.pow(modelGraph[layer] - sess.run(modelGraph[layer]), 2))
  28. # Squared-error loss of style between the two feature representations
  29. styleLayers = [
  30. ('conv1_1', 0.2),
  31. ('conv2_1', 0.2),
  32. ('conv3_1', 0.2),
  33. ('conv4_1', 0.2),
  34. ('conv5_1', 0.2),
  35. ]
  36. def sqErrorLossStyle(sess, modelGraph):
  37. def intermediateCalc(x, y):
  38. N = x.shape[3]
  39. M = x.shape[1] * x.shape[2]
  40. A = tf.matmul(tf.transpose(tf.reshape(x, (M, N))), tf.reshape(x, (M, N)))
  41. G = tf.matmul(tf.transpose(tf.reshape(y, (M, N))), tf.reshape(y, (M, N)))
  42. return (1 / (4 * N**2 * M**2)) * tf.reduce_sum(tf.pow(G - A, 2))
  43. E = [intermediateCalc(sess.run(modelGraph[layerName]), modelGraph[layerName]) for layerName, _ in styleLayers]
  44. W = [w for _, w in styleLayers]
  45. return sum([W[layerNumber] * E[layerNumber] for layerNumber in range(len(styleLayers))])
  46. session = tf.InteractiveSession()
  47. # Addition of extra dimension to image
  48. inputImage = np.reshape(inputImage, ((1,) + inputImage.shape))
  49. inputImage = inputImage - meanImage19
  50. # Display image
  51. pyplot.imshow(inputImage[0])
  52. # Addition of extra dimension to image
  53. paintingStyleImage = np.reshape(paintingStyleImage, ((1,) + paintingStyleImage.shape))
  54. paintingStyleImage = paintingStyleImage - meanImage19
  55. # Display image
  56. pyplot.imshow(paintingStyleImage[0])
  57. imageNoise = np.random.uniform(-20, 20, (1, outputHeight, outputWidth, 3)).astype('float32')
  58. pyplot.imshow(imageNoise[0])
  59. mixedImage = imageNoise * noiseRatio + inputImage * (1 - noiseRatio)
  60. pyplot.imshow(inputImage[0])
  61. session.run(tf.global_variables_initializer())
  62. session.run(nodes['input'].assign(inputImage))
  63. contentLoss = sqErrorLossContent(session, nodes, 'conv4_2')
  64. session.run(nodes['input'].assign(paintingStyleImage))
  65. styleLoss = sqErrorLossStyle(session, nodes)
  66. totalLoss = beta * contentLoss + alpha * styleLoss
  67. optimizer = tf.train.AdamOptimizer(2.0)
  68. trainStep = optimizer.minimize(totalLoss)
  69. session.run(tf.global_variables_initializer())
  70. session.run(nodes['input'].assign(inputImage))
  71. # Number of iterations to run.
  72. iterations = 2000
  73. session.run(tf.global_variables_initializer())
  74. session.run(nodes['input'].assign(inputImage))
  75. for iters in range(iterations):
  76. session.run(trainStep)
  77. if iters%50 == 0:
  78. # Output every 50 iterations for animation
  79. filename = 'output%d.png' % (iters)
  80. im = mixedImage + meanImage19
  81. im = im[0]
  82. im = np.clip(im, 0, 255).astype('uint8')
  83. misc.imsave(filename, im)
  84. im = mixedImage + meanImage19
  85. im = im[0]
  86. im = np.clip(im, 0, 255).astype('uint8')
  87. misc.imsave('finalImage.png', im)