add 'aspect_ratio' in the test code

6 files changed, 36 insertions, 29 deletions

diff --git a/.gitignore b/.gitignore
index f99f990..faba6c9 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,3 +1,4 @@
+.DS_Store
 debug*
 datasets/
 checkpoints/
diff --git a/datasets/combine_A_and_B.py b/datasets/combine_A_and_B.py
index 4d1e2a2..70907ad 100644
--- a/datasets/combine_A_and_B.py
+++ b/datasets/combine_A_and_B.py
@@ -1,4 +1,3 @@
-from pdb import set_trace as st
 import os
 import numpy as np
 import cv2
@@ -21,7 +20,7 @@ for sp in splits:
     img_fold_A = os.path.join(args.fold_A, sp)
     img_fold_B = os.path.join(args.fold_B, sp)
     img_list = os.listdir(img_fold_A)
-    if args.use_AB: 
+    if args.use_AB:
         img_list = [img_path for img_path in img_list if '_A.' in img_path]
 
     num_imgs = min(args.num_imgs, len(img_list))
diff --git a/options/test_options.py b/options/test_options.py
index 7ffc94a..6b79860 100644
--- a/options/test_options.py
+++ b/options/test_options.py
@@ -10,5 +10,4 @@ class TestOptions(BaseOptions):
         self.parser.add_argument('--phase', type=str, default='test', help='train, val, test, etc')
         self.parser.add_argument('--which_epoch', type=str, default='latest', help='which epoch to load? set to latest to use latest cached model')
         self.parser.add_argument('--how_many', type=int, default=50, help='how many test images to run')
-        #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.isTrain = False
diff --git a/test.py b/test.py
index fc0f1bb..65e79d7 100644
--- a/test.py
+++ b/test.py
@@ -28,6 +28,6 @@ for i, data in enumerate(dataset):
     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)
+    visualizer.save_images(webpage, visuals, img_path, aspect_ratio=opt.aspect_ratio)
 
 webpage.save()
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)