Merge pull request #48 from SsnL/visdom_opts

Visdom port and display intermediate in single image pane

2 files changed, 38 insertions, 7 deletions

diff --git a/options/base_options.py b/options/base_options.py
index 44100cf..704aaef 100644
--- a/options/base_options.py
+++ b/options/base_options.py
@@ -32,6 +32,8 @@ class BaseOptions():
         self.parser.add_argument('--serial_batches', action='store_true', help='if true, takes images in order to make batches, otherwise takes them randomly')
         self.parser.add_argument('--display_winsize', type=int, default=256,  help='display window size')
         self.parser.add_argument('--display_id', type=int, default=1, help='window id of the web display')
+        self.parser.add_argument('--display_port', type=int, default=8097, help='visdom port of the web display')
+        self.parser.add_argument('--display_single_pane_ncols', type=int, default=0, help='if positive, display all images in a single visdom web panel with certain number of images per row.')
         self.parser.add_argument('--identity', type=float, default=0.0, help='use identity mapping. Setting identity other than 1 has an effect of scaling the weight of the identity mapping loss. For example, if the weight of the identity loss should be 10 times smaller than the weight of the reconstruction loss, please set optidentity = 0.1')
         self.parser.add_argument('--use_dropout', action='store_true', help='use dropout for the generator')
         self.parser.add_argument('--max_dataset_size', type=int, default=float("inf"), help='Maximum number of samples allowed per dataset. If the dataset directory contains more than max_dataset_size, only a subset is loaded.')
diff --git a/util/visualizer.py b/util/visualizer.py
index a718cfc..38b3bac 100644
--- a/util/visualizer.py
+++ b/util/visualizer.py
@@ -14,7 +14,8 @@ class Visualizer():
         self.name = opt.name
         if self.display_id > 0:
             import visdom
-            self.vis = visdom.Visdom()
+            self.vis = visdom.Visdom(port = opt.display_port)
+            self.display_single_pane_ncols = opt.display_single_pane_ncols
 
         if self.use_html:
             self.web_dir = os.path.join(opt.checkpoints_dir, opt.name, 'web')
@@ -29,12 +30,40 @@ class Visualizer():
     # |visuals|: dictionary of images to display or save
     def display_current_results(self, visuals, epoch):
         if self.display_id > 0: # show images in the browser
-            idx = 1
-            for label, image_numpy in visuals.items():
-                #image_numpy = np.flipud(image_numpy)
-                self.vis.image(image_numpy.transpose([2,0,1]), opts=dict(title=label),
-                                   win=self.display_id + idx)
-                idx += 1
+            if self.display_single_pane_ncols > 0:
+                ncols = self.display_single_pane_ncols
+                title = self.name
+                label_html = ''
+                label_html_row = ''
+                nrows = int(np.ceil(len(visuals.items()) / ncols))
+                images = []
+                idx = 0
+                for label, image_numpy in visuals.items():
+                    label_html_row += '<td>%s</td>' % label
+                    images.append(image_numpy.transpose([2, 0, 1]))
+                    idx += 1
+                    if idx % ncols == 0:
+                        label_html += '<tr>%s</tr>' % label_html_row
+                        label_html_row = ''
+                while idx % ncols != 0:
+                    white_image = np.ones_like(image_numpy.transpose([2, 0, 1]))*255
+                    images.append(white_image)
+                    label_html_row += '<td></td>'
+                    idx += 1
+                if label_html_row != '':
+                    label_html += '<tr>%s</tr>' % label_html_row
+                self.vis.images(images, nrow=ncols, win=self.display_id + 1,
+                              opts=dict(title=title + ' images')) # pane col = image row
+                label_html = '<table style="border-collapse:separate;border-spacing:10px;">%s</table' % label_html
+                self.vis.text(label_html, win = self.display_id + 2,
+                              opts=dict(title=title + ' labels'))
+            else:
+                idx = 1
+                for label, image_numpy in visuals.items():
+                    #image_numpy = np.flipud(image_numpy)
+                    self.vis.image(image_numpy.transpose([2,0,1]), opts=dict(title=label),
+                                       win=self.display_id + idx)
+                    idx += 1
 
         if self.use_html: # save images to a html file
             for label, image_numpy in visuals.items():