All

GANimation/utils/util.py

@@ -0,0 +1,53 @@
+from __future__ import print_function
+from PIL import Image
+import numpy as np
+import os
+import torchvision
+import math
+
+
+def tensor2im(img, imtype=np.uint8, unnormalize=True, idx=0, nrows=None):
+    # select a sample or create grid if img is a batch
+    if len(img.shape) == 4:
+        nrows = nrows if nrows is not None else int(math.sqrt(img.size(0)))
+        img = img[idx] if idx >= 0 else torchvision.utils.make_grid(img, nrows)
+
+    img = img.cpu().float()
+    if unnormalize:
+        mean = [0.5, 0.5, 0.5]
+        std = [0.5, 0.5, 0.5]
+
+        for i, m, s in zip(img, mean, std):
+            i.mul_(s).add_(m)
+
+    image_numpy = img.numpy()
+    image_numpy_t = np.transpose(image_numpy, (1, 2, 0))
+    image_numpy_t = image_numpy_t*254.0
+
+    return image_numpy_t.astype(imtype)
+
+def tensor2maskim(mask, imtype=np.uint8, idx=0, nrows=1):
+    im = tensor2im(mask, imtype=imtype, idx=idx, unnormalize=False, nrows=nrows)
+    if im.shape[2] == 1:
+        im = np.repeat(im, 3, axis=-1)
+    return im
+
+def mkdirs(paths):
+    if isinstance(paths, list) and not isinstance(paths, str):
+        for path in paths:
+            mkdir(path)
+    else:
+        mkdir(paths)
+
+def mkdir(path):
+    if not os.path.exists(path):
+        os.makedirs(path)
+
+def save_image(image_numpy, image_path):
+    mkdir(os.path.dirname(image_path))
+    image_pil = Image.fromarray(image_numpy)
+    image_pil.save(image_path)
+
+def save_str_data(data, path):
+    mkdir(os.path.dirname(path))
+    np.savetxt(path, data, delimiter=",", fmt="%s")