port and single image display

1 files changed, 36 insertions, 7 deletions

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():