This class includes several functions that can display/save images and print/save logging information. It uses a Python library 'visdom' for display, and a Python library 'dominate' (wrapped in 'HTML') for creating HTML files with images.
| 44 | |
| 45 | |
| 46 | class Visualizer(): |
| 47 | """This class includes several functions that can display/save images and print/save logging information. |
| 48 | |
| 49 | It uses a Python library 'visdom' for display, and a Python library 'dominate' (wrapped in 'HTML') for creating HTML files with images. |
| 50 | """ |
| 51 | |
| 52 | def __init__(self, opt): |
| 53 | """Initialize the Visualizer class |
| 54 | |
| 55 | Parameters: |
| 56 | opt -- stores all the experiment flags; needs to be a subclass of BaseOptions |
| 57 | Step 1: Cache the training/test options |
| 58 | Step 2: connect to a visdom server |
| 59 | Step 3: create an HTML object for saveing HTML filters |
| 60 | Step 4: create a logging file to store training losses |
| 61 | """ |
| 62 | self.opt = opt # cache the option |
| 63 | if opt.display_id is None: |
| 64 | self.display_id = np.random.randint(100000) * 10 # just a random display id |
| 65 | else: |
| 66 | self.display_id = opt.display_id |
| 67 | self.use_html = opt.isTrain and not opt.no_html |
| 68 | self.win_size = opt.display_winsize |
| 69 | self.name = opt.name |
| 70 | self.port = opt.display_port |
| 71 | self.saved = False |
| 72 | if self.display_id > 0: # connect to a visdom server given <display_port> and <display_server> |
| 73 | import visdom |
| 74 | self.plot_data = {} |
| 75 | self.ncols = opt.display_ncols |
| 76 | if "tensorboard_base_url" not in os.environ: |
| 77 | self.vis = visdom.Visdom(server=opt.display_server, port=opt.display_port, env=opt.display_env) |
| 78 | else: |
| 79 | self.vis = visdom.Visdom(port=2004, |
| 80 | base_url=os.environ['tensorboard_base_url'] + '/visdom') |
| 81 | if not self.vis.check_connection(): |
| 82 | self.create_visdom_connections() |
| 83 | |
| 84 | if self.use_html: # create an HTML object at <checkpoints_dir>/web/; images will be saved under <checkpoints_dir>/web/images/ |
| 85 | self.web_dir = os.path.join(opt.checkpoints_dir, opt.name, 'web') |
| 86 | self.img_dir = os.path.join(self.web_dir, 'images') |
| 87 | print('create web directory %s...' % self.web_dir) |
| 88 | util.mkdirs([self.web_dir, self.img_dir]) |
| 89 | # create a logging file to store training losses |
| 90 | self.log_name = os.path.join(opt.checkpoints_dir, opt.name, 'loss_log.txt') |
| 91 | with open(self.log_name, "a") as log_file: |
| 92 | now = time.strftime("%c") |
| 93 | log_file.write('================ Training Loss (%s) ================\n' % now) |
| 94 | |
| 95 | def reset(self): |
| 96 | """Reset the self.saved status""" |
| 97 | self.saved = False |
| 98 | |
| 99 | def create_visdom_connections(self): |
| 100 | """If the program could not connect to Visdom server, this function will start a new server at port < self.port > """ |
| 101 | cmd = sys.executable + ' -m visdom.server -p %d &>/dev/null &' % self.port |
| 102 | print('\n\nCould not connect to Visdom server. \n Trying to start a server....') |
| 103 | print('Command: %s' % cmd) |