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