add 'aspect_ratio' in the test code

2 files changed, 33 insertions, 25 deletions

diff --git a/util/png.py b/util/png.py
index 0936cf0..3a750c0 100644
--- a/util/png.py
+++ b/util/png.py
@@ -1,32 +1,33 @@
 import struct
 import zlib
 
+
 def encode(buf, width, height):
-  """ buf: must be bytes or a bytearray in py3, a regular string in py2. formatted RGBRGB... """
-  assert (width * height * 3 == len(buf))
-  bpp = 3
+    """ buf: must be bytes or a bytearray in py3, a regular string in py2. formatted RGBRGB... """
+    assert (width * height * 3 == len(buf))
+    bpp = 3
 
-  def raw_data():
-    # reverse the vertical line order and add null bytes at the start
-    row_bytes = width * bpp
-    for row_start in range((height - 1) * width * bpp, -1, -row_bytes):
-      yield b'\x00'
-      yield buf[row_start:row_start + row_bytes]
+    def raw_data():
+        # reverse the vertical line order and add null bytes at the start
+        row_bytes = width * bpp
+        for row_start in range((height - 1) * width * bpp, -1, -row_bytes):
+            yield b'\x00'
+            yield buf[row_start:row_start + row_bytes]
 
-  def chunk(tag, data):
-    return [
-        struct.pack("!I", len(data)),
-        tag,
-        data,
-        struct.pack("!I", 0xFFFFFFFF & zlib.crc32(data, zlib.crc32(tag)))
-      ]
+    def chunk(tag, data):
+        return [
+            struct.pack("!I", len(data)),
+            tag,
+            data,
+            struct.pack("!I", 0xFFFFFFFF & zlib.crc32(data, zlib.crc32(tag)))
+            ]
 
-  SIGNATURE = b'\x89PNG\r\n\x1a\n'
-  COLOR_TYPE_RGB = 2
-  COLOR_TYPE_RGBA = 6
-  bit_depth = 8
-  return b''.join(
-      [ SIGNATURE ] +
+    SIGNATURE = b'\x89PNG\r\n\x1a\n'
+    COLOR_TYPE_RGB = 2
+    COLOR_TYPE_RGBA = 6
+    bit_depth = 8
+    return b''.join(
+      [SIGNATURE] +
       chunk(b'IHDR', struct.pack("!2I5B", width, height, bit_depth, COLOR_TYPE_RGB, 0, 0, 0)) +
       chunk(b'IDAT', zlib.compress(b''.join(raw_data()), 9)) +
       chunk(b'IEND', b'')
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)