All

GANimation/utils/cv_utils.py

@@ -0,0 +1,54 @@
+import cv2
+from matplotlib import pyplot as plt
+import numpy as np
+
+def read_cv2_img(path):
+    '''
+    Read color images
+    :param path: Path to image
+    :return: Only returns color images
+    '''
+    img = cv2.imread(path, -1)
+
+    if img is not None:
+        if len(img.shape) != 3:
+            return None
+
+        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+
+    return img
+
+def show_cv2_img(img, title='img'):
+    '''
+    Display cv2 image
+    :param img: cv::mat
+    :param title: title
+    :return: None
+    '''
+    plt.imshow(img)
+    plt.title(title)
+    plt.axis('off')
+    plt.show()
+
+def show_images_row(imgs, titles, rows=1):
+    '''
+       Display grid of cv2 images image
+       :param img: list [cv::mat]
+       :param title: titles
+       :return: None
+    '''
+    assert ((titles is None) or (len(imgs) == len(titles)))
+    num_images = len(imgs)
+
+    if titles is None:
+        titles = ['Image (%d)' % i for i in range(1, num_images + 1)]
+
+    fig = plt.figure()
+    for n, (image, title) in enumerate(zip(imgs, titles)):
+        ax = fig.add_subplot(rows, np.ceil(num_images / float(rows)), n + 1)
+        if image.ndim == 2:
+            plt.gray()
+        plt.imshow(image)
+        ax.set_title(title)
+        plt.axis('off')
+    plt.show()

GANimation/utils/face_utils.py

@@ -0,0 +1,71 @@
+import face_recognition
+import cv2
+import numpy as np
+import skimage
+import skimage.transform
+import warnings
+
+def detect_faces(img):
+    '''
+    Detect faces in image
+    :param img: cv::mat HxWx3 RGB
+    :return: yield 4 <x,y,w,h>
+    '''
+    # detect faces
+    bbs = face_recognition.face_locations(img)
+
+    for y, right, bottom, x in bbs:
+        # Scale back up face bb
+        yield x, y, (right - x), (bottom - y)
+
+def detect_biggest_face(img):
+    '''
+    Detect biggest face in image
+    :param img: cv::mat HxWx3 RGB
+    :return: 4 <x,y,w,h>
+    '''
+    # detect faces
+    bbs = face_recognition.face_locations(img)
+
+    max_area = float('-inf')
+    max_area_i = 0
+    for i, (y, right, bottom, x) in enumerate(bbs):
+        area = (right - x) * (bottom - y)
+        if max_area < area:
+            max_area = area
+            max_area_i = i
+
+    if max_area != float('-inf'):
+        y, right, bottom, x = bbs[max_area_i]
+        return x, y, (right - x), (bottom - y)
+
+    return None
+
+def crop_face_with_bb(img, bb):
+    '''
+    Crop face in image given bb
+    :param img: cv::mat HxWx3
+    :param bb: 4 (<x,y,w,h>)
+    :return: HxWx3
+    '''
+    x, y, w, h = bb
+    return img[y:y+h, x:x+w, :]
+
+def place_face(img, face, bb):
+    x, y, w, h = bb
+    face = resize_face(face, size=(w, h))
+    img[y:y+h, x:x+w] = face
+    return img
+
+def resize_face(face_img, size=(128, 128)):
+    '''
+    Resize face to a given size
+    :param face_img: cv::mat HxWx3
+    :param size: new H and W (size x size). 128 by default.
+    :return: cv::mat size x size x 3
+    '''
+    return cv2.resize(face_img, size)
+
+def detect_landmarks(face_img):
+    landmakrs = face_recognition.face_landmarks(face_img)
+    return landmakrs[0] if len(landmakrs) > 0 else None

GANimation/utils/plots.py

@@ -0,0 +1,67 @@
+from __future__ import print_function
+import numpy as np
+import matplotlib.pyplot as plt
+
+def plot_au(img, aus, title=None):
+    '''
+    Plot action units
+    :param img: HxWx3
+    :param aus: N
+    :return:
+    '''
+    fig = plt.figure()
+    ax = fig.add_subplot(1, 1, 1)
+    ax.axis('off')
+    fig.subplots_adjust(0, 0, 0.8, 1)  # get rid of margins
+
+    # display img
+    ax.imshow(img)
+
+    if len(aus) == 11:
+        au_ids = ['1','2','4','5','6','9','12','17','20','25','26']
+        x = 0.1
+        y = 0.39
+        i = 0
+        for au, id in zip(aus, au_ids):
+            if id == '9':
+                x = 0.5
+                y -= .15
+                i = 0
+            elif id == '12':
+                x = 0.1
+                y -= .15
+                i = 0
+
+            ax.text(x + i * 0.2, y, id, horizontalalignment='center', verticalalignment='center',
+                    transform=ax.transAxes, color='r', fontsize=20)
+            ax.text((x-0.001)+i*0.2, y-0.07, au, horizontalalignment='center', verticalalignment='center',
+                    transform=ax.transAxes, color='b', fontsize=20)
+            i+=1
+
+    else:
+        au_ids = ['1', '2', '4', '5', '6', '7', '9', '10', '12', '14', '15', '17', '20', '23', '25', '26', '45']
+        x = 0.1
+        y = 0.39
+        i = 0
+        for au, id in zip(aus, au_ids):
+            if id == '9' or id == '20':
+                x = 0.1
+                y -= .15
+                i = 0
+
+            ax.text(x + i * 0.2, y, id, horizontalalignment='center', verticalalignment='center',
+                    transform=ax.transAxes, color='r', fontsize=20)
+            ax.text((x-0.001)+i*0.2, y-0.07, au, horizontalalignment='center', verticalalignment='center',
+                    transform=ax.transAxes, color='b', fontsize=20)
+            i+=1
+
+    if title is not None:
+        ax.text(0.5, 0.95, title, horizontalalignment='center', verticalalignment='center',
+                transform=ax.transAxes, color='r', fontsize=20)
+
+    fig.canvas.draw()
+    data = np.fromstring(fig.canvas.tostring_rgb(), dtype=np.uint8, sep='')
+    data = data.reshape(fig.canvas.get_width_height()[::-1] + (3,))
+    plt.close(fig)
+
+    return data

GANimation/utils/tb_visualizer.py

@@ -0,0 +1,66 @@
+import numpy as np
+import os
+import time
+from . import util
+from tensorboardX import SummaryWriter
+
+
+class TBVisualizer:
+    def __init__(self, opt):
+        self._opt = opt
+        self._save_path = os.path.join(opt.checkpoints_dir, opt.name)
+
+        self._log_path = os.path.join(self._save_path, 'loss_log2.txt')
+        self._tb_path = os.path.join(self._save_path, 'summary.json')
+        self._writer = SummaryWriter(self._save_path)
+
+        with open(self._log_path, "a") as log_file:
+            now = time.strftime("%c")
+            log_file.write('================ Training Loss (%s) ================\n' % now)
+
+    def __del__(self):
+        self._writer.close()
+
+    def display_current_results(self, visuals, it, is_train, save_visuals=False):
+        for label, image_numpy in visuals.items():
+            sum_name = '{}/{}'.format('Train' if is_train else 'Test', label)
+            self._writer.add_image(sum_name, image_numpy, it)
+
+            if save_visuals:
+                util.save_image(image_numpy,
+                                os.path.join(self._opt.checkpoints_dir, self._opt.name,
+                                             'event_imgs', sum_name, '%08d.png' % it))
+
+        self._writer.export_scalars_to_json(self._tb_path)
+
+    def plot_scalars(self, scalars, it, is_train):
+        for label, scalar in scalars.items():
+            sum_name = '{}/{}'.format('Train' if is_train else 'Test', label)
+            self._writer.add_scalar(sum_name, scalar, it)
+
+    def print_current_train_errors(self, epoch, i, iters_per_epoch, errors, t, visuals_were_stored):
+        log_time = time.strftime("[%d/%m/%Y %H:%M:%S]")
+        visuals_info = "v" if visuals_were_stored else ""
+        message = '%s (T%s, epoch: %d, it: %d/%d, t/smpl: %.3fs) ' % (log_time, visuals_info, epoch, i, iters_per_epoch, t)
+        for k, v in errors.items():
+            message += '%s:%.3f ' % (k, v)
+
+        print(message)
+        with open(self._log_path, "a") as log_file:
+            log_file.write('%s\n' % message)
+
+    def print_current_validate_errors(self, epoch, errors, t):
+        log_time = time.strftime("[%d/%m/%Y %H:%M:%S]")
+        message = '%s (V, epoch: %d, time_to_val: %ds) ' % (log_time, epoch, t)
+        for k, v in errors.items():
+            message += '%s:%.3f ' % (k, v)
+
+        print(message)
+        with open(self._log_path, "a") as log_file:
+            log_file.write('%s\n' % message)
+
+    def save_images(self, visuals):
+        for label, image_numpy in visuals.items():
+            image_name = '%s.png' % label
+            save_path = os.path.join(self._save_path, "samples", image_name)
+            util.save_image(image_numpy, save_path)

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")