added image save callback to sample/render loop

1 files changed, 19 insertions, 0 deletions

diff --git a/gantools/image_utils.py b/gantools/image_utils.py
index 72130b0..a875a91 100644
--- a/gantools/image_utils.py
+++ b/gantools/image_utils.py
@@ -9,3 +9,22 @@ def save_images(ims, output_dir='', prefix='', format='JPEG'):
     for i, im in enumerate(ims):
         full_path = os.path.join(output_dir, prefix + str(i).zfill(4) + '.' + format.lower())
         save_image(im, full_path, format)
+
+class ImageSaver(object):
+    def __init__(self, output_dir='', prefix='', image_format='JPEG'):
+        self.output_dir = str(output_dir)
+        self.prefix = str(prefix)
+        self.image_format = str(image_format)
+        self.index = int(0)
+        self.quality = 90
+
+    def save(self, ims):
+        for i, im in enumerate(ims):
+            full_path = os.path.join(
+                    self.output_dir,
+                    self.prefix + str(self.index).zfill(4) + '.' + self.image_format.lower()
+                    )
+            image = PIL.Image.fromarray(im)
+            image.save(full_path, format=self.image_format, quality=self.quality)
+            self.index += 1
+