simplify visualizer

Author: junyanz

test.py

@@ -2,7 +2,7 @@
 from options.test_options import TestOptions
 from data import CreateDataLoader
 from models import create_model
-from util.visualizer import Visualizer
+from util.visualizer import save_images
 from util import html
 
 
@@ -16,7 +16,6 @@
     data_loader = CreateDataLoader(opt)
     dataset = data_loader.load_data()
     model = create_model(opt)
-    visualizer = Visualizer(opt)
     # create website
     web_dir = os.path.join(opt.results_dir, opt.name, '%s_%s' % (opt.phase, opt.which_epoch))
     webpage = html.HTML(web_dir, 'Experiment = %s, Phase = %s, Epoch = %s' % (opt.name, opt.phase, opt.which_epoch))
@@ -28,7 +27,8 @@
         model.test()
         visuals = model.get_current_visuals()
         img_path = model.get_image_paths()
-        print('%04d: process image... %s' % (i, img_path))
-        visualizer.save_images(webpage, visuals, img_path, aspect_ratio=opt.aspect_ratio)
+        if i % 5 == 0:
+            print('processing (%04d)-th image... %s' % (i, img_path))
+        save_images(webpage, visuals, img_path, aspect_ratio=opt.aspect_ratio, width=opt.display_winsize)
 
     webpage.save()

util/visualizer.py

@@ -7,6 +7,32 @@
 from scipy.misc import imresize
 
 
+# save image to the disk
+def save_images(webpage, visuals, image_path, aspect_ratio=1.0, width=256):
+    image_dir = webpage.get_image_dir()
+    short_path = ntpath.basename(image_path[0])
+    name = os.path.splitext(short_path)[0]
+
+    webpage.add_header(name)
+    ims, txts, links = [], [], []
+
+    for label, im_data in visuals.items():
+        im = util.tensor2im(im_data)
+        image_name = '%s_%s.png' % (name, label)
+        save_path = os.path.join(image_dir, image_name)
+        h, w, _ = im.shape
+        if aspect_ratio > 1.0:
+            im = imresize(im, (h, int(w * aspect_ratio)), interp='bicubic')
+        if aspect_ratio < 1.0:
+            im = imresize(im, (int(h / aspect_ratio), w), interp='bicubic')
+        util.save_image(im, save_path)
+
+        ims.append(image_name)
+        txts.append(label)
+        links.append(image_name)
+    webpage.add_images(ims, txts, links, width=width)
+
+
 class Visualizer():
     def __init__(self, opt):
         self.display_id = opt.display_id
@@ -124,28 +150,3 @@ def print_current_losses(self, epoch, i, losses, t, t_data):
         print(message)
         with open(self.log_name, "a") as log_file:
             log_file.write('%s\n' % message)
-
-    # save image to the disk
-    def save_images(self, webpage, visuals, image_path, aspect_ratio=1.0):
-        image_dir = webpage.get_image_dir()
-        short_path = ntpath.basename(image_path[0])
-        name = os.path.splitext(short_path)[0]
-
-        webpage.add_header(name)
-        ims, txts, links = [], [], []
-
-        for label, im_data in visuals.items():
-            im = util.tensor2im(im_data)
-            image_name = '%s_%s.png' % (name, label)
-            save_path = os.path.join(image_dir, image_name)
-            h, w, _ = im.shape
-            if aspect_ratio > 1.0:
-                im = imresize(im, (h, int(w * aspect_ratio)), interp='bicubic')
-            if aspect_ratio < 1.0:
-                im = imresize(im, (int(h / aspect_ratio), w), interp='bicubic')
-            util.save_image(im, save_path)
-
-            ims.append(image_name)
-            txts.append(label)
-            links.append(image_name)
-        webpage.add_images(ims, txts, links, width=self.win_size)