visualizer.py 5.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131
  1. import numpy as np
  2. import os
  3. import ntpath
  4. import time
  5. from . import util
  6. from . import html
  7. import scipy.misc
  8. try:
  9. from StringIO import StringIO # Python 2.7
  10. except ImportError:
  11. from io import BytesIO # Python 3.x
  12. class Visualizer():
  13. def __init__(self, opt):
  14. # self.opt = opt
  15. self.tf_log = opt.tf_log
  16. self.use_html = opt.isTrain and not opt.no_html
  17. self.win_size = opt.display_winsize
  18. self.name = opt.name
  19. if self.tf_log:
  20. import tensorflow as tf
  21. self.tf = tf
  22. self.log_dir = os.path.join(opt.checkpoints_dir, opt.name, 'logs')
  23. self.writer = tf.summary.FileWriter(self.log_dir)
  24. if self.use_html:
  25. self.web_dir = os.path.join(opt.checkpoints_dir, opt.name, 'web')
  26. self.img_dir = os.path.join(self.web_dir, 'images')
  27. print('create web directory %s...' % self.web_dir)
  28. util.mkdirs([self.web_dir, self.img_dir])
  29. self.log_name = os.path.join(opt.checkpoints_dir, opt.name, 'loss_log.txt')
  30. with open(self.log_name, "a") as log_file:
  31. now = time.strftime("%c")
  32. log_file.write('================ Training Loss (%s) ================\n' % now)
  33. # |visuals|: dictionary of images to display or save
  34. def display_current_results(self, visuals, epoch, step):
  35. if self.tf_log: # show images in tensorboard output
  36. img_summaries = []
  37. for label, image_numpy in visuals.items():
  38. # Write the image to a string
  39. try:
  40. s = StringIO()
  41. except:
  42. s = BytesIO()
  43. scipy.misc.toimage(image_numpy).save(s, format="jpeg")
  44. # Create an Image object
  45. img_sum = self.tf.Summary.Image(encoded_image_string=s.getvalue(), height=image_numpy.shape[0], width=image_numpy.shape[1])
  46. # Create a Summary value
  47. img_summaries.append(self.tf.Summary.Value(tag=label, image=img_sum))
  48. # Create and write Summary
  49. summary = self.tf.Summary(value=img_summaries)
  50. self.writer.add_summary(summary, step)
  51. if self.use_html: # save images to a html file
  52. for label, image_numpy in visuals.items():
  53. if isinstance(image_numpy, list):
  54. for i in range(len(image_numpy)):
  55. img_path = os.path.join(self.img_dir, 'epoch%.3d_%s_%d.jpg' % (epoch, label, i))
  56. util.save_image(image_numpy[i], img_path)
  57. else:
  58. img_path = os.path.join(self.img_dir, 'epoch%.3d_%s.jpg' % (epoch, label))
  59. util.save_image(image_numpy, img_path)
  60. # update website
  61. webpage = html.HTML(self.web_dir, 'Experiment name = %s' % self.name, refresh=30)
  62. for n in range(epoch, 0, -1):
  63. webpage.add_header('epoch [%d]' % n)
  64. ims = []
  65. txts = []
  66. links = []
  67. for label, image_numpy in visuals.items():
  68. if isinstance(image_numpy, list):
  69. for i in range(len(image_numpy)):
  70. img_path = 'epoch%.3d_%s_%d.jpg' % (n, label, i)
  71. ims.append(img_path)
  72. txts.append(label+str(i))
  73. links.append(img_path)
  74. else:
  75. img_path = 'epoch%.3d_%s.jpg' % (n, label)
  76. ims.append(img_path)
  77. txts.append(label)
  78. links.append(img_path)
  79. if len(ims) < 10:
  80. webpage.add_images(ims, txts, links, width=self.win_size)
  81. else:
  82. num = int(round(len(ims)/2.0))
  83. webpage.add_images(ims[:num], txts[:num], links[:num], width=self.win_size)
  84. webpage.add_images(ims[num:], txts[num:], links[num:], width=self.win_size)
  85. webpage.save()
  86. # errors: dictionary of error labels and values
  87. def plot_current_errors(self, errors, step):
  88. if self.tf_log:
  89. for tag, value in errors.items():
  90. summary = self.tf.Summary(value=[self.tf.Summary.Value(tag=tag, simple_value=value)])
  91. self.writer.add_summary(summary, step)
  92. # errors: same format as |errors| of plotCurrentErrors
  93. def print_current_errors(self, epoch, i, errors, t):
  94. message = '(epoch: %d, iters: %d, time: %.3f) ' % (epoch, i, t)
  95. for k, v in errors.items():
  96. if v != 0:
  97. message += '%s: %.3f ' % (k, v)
  98. print(message)
  99. with open(self.log_name, "a") as log_file:
  100. log_file.write('%s\n' % message)
  101. # save image to the disk
  102. def save_images(self, webpage, visuals, image_path):
  103. image_dir = webpage.get_image_dir()
  104. short_path = ntpath.basename(image_path[0])
  105. name = os.path.splitext(short_path)[0]
  106. webpage.add_header(name)
  107. ims = []
  108. txts = []
  109. links = []
  110. for label, image_numpy in visuals.items():
  111. image_name = '%s_%s.jpg' % (name, label)
  112. save_path = os.path.join(image_dir, image_name)
  113. util.save_image(image_numpy, save_path)
  114. ims.append(image_name)
  115. txts.append(label)
  116. links.append(image_name)
  117. webpage.add_images(ims, txts, links, width=self.win_size)