| 12 | from io import BytesIO # Python 3.x |
| 13 | |
| 14 | class Visualizer(): |
| 15 | def __init__(self, opt): |
| 16 | self.opt = opt |
| 17 | self.tf_log = opt.tf_log |
| 18 | self.use_html = opt.isTrain and not opt.no_html |
| 19 | self.win_size = opt.display_winsize |
| 20 | self.name = opt.name |
| 21 | if self.tf_log: |
| 22 | import tensorflow as tf |
| 23 | self.tf = tf |
| 24 | self.log_dir = os.path.join(opt.checkpoints_dir, opt.name, 'logs') |
| 25 | self.writer = tf.summary.FileWriter(self.log_dir) |
| 26 | |
| 27 | if self.use_html: |
| 28 | self.web_dir = os.path.join(opt.checkpoints_dir, opt.name, 'web') |
| 29 | self.img_dir = os.path.join(self.web_dir, 'images') |
| 30 | print('create web directory %s...' % self.web_dir) |
| 31 | util.mkdirs([self.web_dir, self.img_dir]) |
| 32 | self.log_name = os.path.join(opt.checkpoints_dir, opt.name, 'loss_log.txt') |
| 33 | with open(self.log_name, "a") as log_file: |
| 34 | now = time.strftime("%c") |
| 35 | log_file.write('================ Training Loss (%s) ================\n' % now) |
| 36 | |
| 37 | # |visuals|: dictionary of images to display or save |
| 38 | def display_current_results(self, visuals, epoch, step): |
| 39 | if self.tf_log: # show images in tensorboard output |
| 40 | img_summaries = [] |
| 41 | for label, image_numpy in visuals.items(): |
| 42 | # Write the image to a string |
| 43 | try: |
| 44 | s = StringIO() |
| 45 | except: |
| 46 | s = BytesIO() |
| 47 | scipy.misc.toimage(image_numpy).save(s, format="jpeg") |
| 48 | # Create an Image object |
| 49 | img_sum = self.tf.Summary.Image(encoded_image_string=s.getvalue(), height=image_numpy.shape[0], width=image_numpy.shape[1]) |
| 50 | # Create a Summary value |
| 51 | img_summaries.append(self.tf.Summary.Value(tag=label, image=img_sum)) |
| 52 | |
| 53 | # Create and write Summary |
| 54 | summary = self.tf.Summary(value=img_summaries) |
| 55 | self.writer.add_summary(summary, step) |
| 56 | |
| 57 | if self.use_html: # save images to a html file |
| 58 | for label, image_numpy in visuals.items(): |
| 59 | if isinstance(image_numpy, list): |
| 60 | for i in range(len(image_numpy)): |
| 61 | img_path = os.path.join(self.img_dir, 'epoch%.3d_%s_%d.jpg' % (epoch, label, i)) |
| 62 | util.save_image(image_numpy[i], img_path) |
| 63 | else: |
| 64 | img_path = os.path.join(self.img_dir, 'epoch%.3d_%s.jpg' % (epoch, label)) |
| 65 | util.save_image(image_numpy, img_path) |
| 66 | |
| 67 | # update website |
| 68 | webpage = html.HTML(self.web_dir, 'Experiment name = %s' % self.name, reflesh=1) |
| 69 | for n in range(epoch, 0, -1): |
| 70 | webpage.add_header('epoch [%d]' % n) |
| 71 | ims = [] |