add 'aspect_ratio' in the test code

1 files changed, 10 insertions, 3 deletions

diff --git a/util/visualizer.py b/util/visualizer.py
index e6e7cba..58d1a2a 100644
--- a/util/visualizer.py
+++ b/util/visualizer.py
@@ -4,6 +4,8 @@ import ntpath
 import time
 from . import util
 from . import html
+from scipy.misc import imresize
+from pdb import set_trace as st
 
 
 class Visualizer():
@@ -123,7 +125,7 @@ class Visualizer():
             log_file.write('%s\n' % message)
 
     # save image to the disk
-    def save_images(self, webpage, visuals, image_path):
+    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]
@@ -133,10 +135,15 @@ class Visualizer():
         txts = []
         links = []
 
-        for label, image_numpy in visuals.items():
+        for label, im in visuals.items():
             image_name = '%s_%s.png' % (name, label)
             save_path = os.path.join(image_dir, image_name)
-            util.save_image(image_numpy, save_path)
+            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)