dfl mask. Make masks selectable in config (#575)

* DFL H128 Mask. Mask type selectable in config.

lib/config.py

@@ -203,6 +203,7 @@ def validate_config(self):
                                         opt["helptext"],
                                         new_config)
         self.config = new_config
+        self.config.optionxform = str
         self.save_config()
         logger.debug("Updated config")
 

lib/model/masks.py

@@ -3,15 +3,22 @@
     Masks from:
         dfaker: https://github.com/dfaker/df"""
 
+import logging
+
 import cv2
 import numpy as np
 
 from lib.umeyama import umeyama
 from lib.aligner import LANDMARKS_2D
 
+logger = logging.getLogger(__name__)  # pylint: disable=invalid-name
+
 
-def dfaker_mask(landmarks, coverage, face):
-    """ Dfaker model mask """
+def dfaker(landmarks, face, **kwargs):
+    """ Dfaker model mask
+        Embeds the mask into the face alpha channel """
+    coverage = kwargs["coverage"]
+    logger.trace("face_shape: %s, coverage: %s, landmarks: %s", face.shape, coverage, landmarks)
     size = face.shape[0] - 1
 
     mat = umeyama(landmarks[17:], LANDMARKS_2D, True)[0:2]
@@ -37,4 +44,31 @@ def dfaker_mask(landmarks, coverage, face):
     mask = mask[:, :, 0]
     merged = np.dstack([face, mask]).astype(np.uint8)
 
+    logger.trace("Returning: face_shape: %s", merged.shape)
     return merged
+
+
+def dfl_full(landmarks, face, **kwargs):
+    """ DFL Face Full Mask """
+    logger.trace("face_shape: %s, landmarks: %s", face.shape, landmarks)
+    hull_mask = np.zeros(face.shape[0:2] + (1, ), dtype=np.float32)
+    hull1 = cv2.convexHull(np.concatenate((landmarks[0:17],  # pylint: disable=no-member
+                                           landmarks[48:],
+                                           [landmarks[0]],
+                                           [landmarks[8]],
+                                           [landmarks[16]])))
+    hull2 = cv2.convexHull(np.concatenate((landmarks[27:31],  # pylint: disable=no-member
+                                           [landmarks[33]])))
+    hull3 = cv2.convexHull(np.concatenate((landmarks[17:27],  # pylint: disable=no-member
+                                           [landmarks[0]],
+                                           [landmarks[27]],
+                                           [landmarks[16]],
+                                           [landmarks[33]])))
+
+    cv2.fillConvexPoly(hull_mask, hull1, (1, ))  # pylint: disable=no-member
+    cv2.fillConvexPoly(hull_mask, hull2, (1, ))  # pylint: disable=no-member
+    cv2.fillConvexPoly(hull_mask, hull3, (1, ))  # pylint: disable=no-member
+
+    face = np.concatenate((face, hull_mask), -1)
+    logger.trace("Returning: face_shape: %s", face.shape)
+    return face

lib/training_data.py

@@ -11,13 +11,16 @@
 import numpy as np
 from scipy.interpolate import griddata
 
+from lib.model import masks
 from lib.multithreading import MultiThread
 from lib.queue_manager import queue_manager
 from lib.umeyama import umeyama
-from lib.model.masks import dfaker_mask
 
 logger = logging.getLogger(__name__)  # pylint: disable=invalid-name
 
+# TODO Add source, dest points to random warp half and ability to
+# not have landmarks to random warp full
+
 
 class TrainingDataGenerator():
     """ Generate training data for models """
@@ -27,9 +30,22 @@ def __init__(self, transform_kwargs, training_opts=dict()):
                      {key: val for key, val in training_opts.items() if key != "landmarks"})
         self.batchsize = 0
         self.training_opts = training_opts
+        self.full_face = self.training_opts.get("full_face", False)
+        self.mask_function = self.set_mask_function()
         self.processing = ImageManipulation(**transform_kwargs)
         logger.debug("Initialized %s", self.__class__.__name__)
 
+    def set_mask_function(self):
+        """ Set the mask function to use if using mask """
+        mask_type = self.training_opts.get("mask_type", None)
+        if mask_type:
+            logger.debug("Mask type: '%s'", mask_type)
+            mask_func = getattr(masks, mask_type)
+        else:
+            mask_func = None
+        logger.debug("Mask function: %s", mask_func)
+        return mask_func
+
     def minibatch_ab(self, images, batchsize, side, do_shuffle=True):
         """ Keep a queue filled to 8x Batch Size """
         logger.debug("Queue batches: (image_count: %s, batchsize: %s, side: '%s', do_shuffle: %s)",
@@ -94,32 +110,29 @@ def minibatch(self, q_name, load_thread):
 
     def process_face(self, filename, side):
         """ Load an image and perform transformation and warping """
-        logger.trace("Process face: (filename: '%s', side: '%s'", filename, side)
+        logger.trace("Process face: (filename: '%s', side: '%s')", filename, side)
         try:
             image = cv2.imread(filename)  # pylint: disable=no-member
         except TypeError:
             raise Exception("Error while reading image", filename)
 
-        landmarks = self.training_opts.get("landmarks", None)
-        src_pts = self.get_landmarks(filename,
-                                     image,
-                                     side,
-                                     landmarks) if landmarks else None
-
-        if self.training_opts.get("use_mask", False):
-            image = dfaker_mask(src_pts, self.processing.coverage, image)
+        if self.mask_function:
+            landmarks = self.training_opts["landmarks"]
+            src_pts = self.get_landmarks(filename, image, side, landmarks)
+            image = self.mask_function(src_pts, image, coverage=self.processing.coverage)
 
         image = self.processing.color_adjust(image)
         image = self.processing.random_transform(image)
 
-        if landmarks:
+        if self.full_face:
             dst_pts = self.get_closest_match(filename, side, landmarks, src_pts)
-            processed = self.processing.random_warp_landmarks(image, src_pts, dst_pts)
+            processed = self.processing.random_warp_full_face(image, src_pts, dst_pts)
         else:
             processed = self.processing.random_warp(image)
 
         retval = self.processing.do_random_flip(processed)
-        logger.trace("Processed face: (filename: '%s', side: '%s'", filename, side)
+        logger.trace("Processed face: (filename: '%s', side: '%s', shapes: %s)",
+                     filename, side, [img.shape for img in retval])
         return retval
 
     @staticmethod
@@ -153,7 +166,7 @@ def get_closest_match(filename, side, landmarks, src_points):
 class ImageManipulation():
     """ Manipulations to be performed on training images """
     def __init__(self, rotation_range=10, zoom_range=0.05, shift_range=0.05, random_flip=0.4,
-                 zoom=1, coverage=160, scale=5):
+                 zoom=(1, 1), coverage=160, scale=5):
         """ rotation_range: Used for random transform
             zoom_range: Used for random transform
             shift_range: Used for random transform
@@ -184,6 +197,16 @@ def color_adjust(img):
         logger.trace("Color adjusting image")
         return img / 255.0
 
+    @staticmethod
+    def seperate_mask(image):
+        """ Return the image and the mask from a 4 channel image """
+        mask = None
+        if image.shape[2] == 4:
+            logger.trace("Image contains mask")
+            mask = image[:, :, 3].reshape((image.shape[0], image.shape[1], 1))
+            image = image[:, :, :3]
+        return image, mask
+
     def random_transform(self, image):
         """ Randomly transform an image """
         logger.trace("Randomly transforming image")
@@ -204,9 +227,10 @@ def random_transform(self, image):
         logger.trace("Randomly transformed image")
         return result
 
-    def random_warp(self, image):
+    def random_warp(self, image, src_points=None, dst_points=None):
         """ get pair of random warped images from aligned face image """
         logger.trace("Randomly warping image")
+        image, mask = self.seperate_mask(image)
         height, width = image.shape[0:2]
         assert height == width and height % 2 == 0
 
@@ -219,31 +243,45 @@ def random_warp(self, image):
         mapy = mapy + np.random.normal(size=(self.scale, self.scale), scale=self.scale)
 
         interp_mapx = cv2.resize(  # pylint: disable=no-member
-            mapx, (80 * self.zoom, 80 * self.zoom))[8 * self.zoom:72 * self.zoom,
-                                                    8 * self.zoom:72 * self.zoom].astype('float32')
+            mapx, (80 * self.zoom[0],
+                   80 * self.zoom[0]))[8 * self.zoom[0]:72 * self.zoom[0],
+                                       8 * self.zoom[0]:72 * self.zoom[0]].astype('float32')
         interp_mapy = cv2.resize(  # pylint: disable=no-member
-            mapy, (80 * self.zoom, 80 * self.zoom))[8 * self.zoom:72 * self.zoom,
-                                                    8 * self.zoom:72 * self.zoom].astype('float32')
+            mapy, (80 * self.zoom[0],
+                   80 * self.zoom[0]))[8 * self.zoom[0]:72 * self.zoom[0],
+                                       8 * self.zoom[0]:72 * self.zoom[0]].astype('float32')
 
         warped_image = cv2.remap(  # pylint: disable=no-member
             image, interp_mapx, interp_mapy, cv2.INTER_LINEAR)  # pylint: disable=no-member
+        logger.trace("Warped image shape: %s", warped_image.shape)
 
         src_points = np.stack([mapx.ravel(), mapy.ravel()], axis=-1)
-        dst_points = np.mgrid[0:65 * self.zoom:16 * self.zoom,
-                              0:65 * self.zoom:16 * self.zoom].T.reshape(-1, 2)
+        dst_points = np.mgrid[0:65 * self.zoom[0]:16 * self.zoom[0],
+                              0:65 * self.zoom[0]:16 * self.zoom[0]].T.reshape(-1, 2)
 
         mat = umeyama(src_points, dst_points, True)[0:2]
         target_image = cv2.warpAffine(  # pylint: disable=no-member
-            image, mat, (64 * self.zoom, 64 * self.zoom))
+            image, mat, (64 * self.zoom[1], 64 * self.zoom[1]))
+        logger.trace("Target image shape: %s", target_image.shape)
+
+        retval = [warped_image, target_image]
+
+        if mask is not None:
+            target_mask = cv2.warpAffine(  # pylint: disable=no-member
+                mask, mat, (64 * self.zoom[1], 64 * self.zoom[1]))
+            target_mask = target_mask.reshape((64 * self.zoom[1], 64 * self.zoom[1], 1))
+            logger.trace("Target mask shape: %s", target_mask.shape)
+
+            retval.append(target_mask)
 
-        retval = warped_image, target_image
         logger.trace("Randomly warped image")
         return retval
 
-    def random_warp_landmarks(self, image, src_points, dst_points):
+    def random_warp_full_face(self, image, src_points=None, dst_points=None):
         """ get warped image, target image and target mask
             From DFAKER plugin """
         logger.trace("Randomly warping landmarks")
+        image, mask = self.seperate_mask(image)
         size = image.shape[0]
         p_mx = size - 1
         p_hf = (size // 2) - 1
@@ -286,36 +324,39 @@ def random_warp_landmarks(self, image, src_points, dst_points):
         map_x_32 = map_x.astype('float32')
         map_y_32 = map_y.astype('float32')
 
-        warped_image = cv2.remap(image[:, :, :3],  # pylint: disable=no-member
+        warped_image = cv2.remap(image,  # pylint: disable=no-member
                                  map_x_32,
                                  map_y_32,
                                  cv2.INTER_LINEAR,  # pylint: disable=no-member
                                  cv2.BORDER_TRANSPARENT)  # pylint: disable=no-member
-        target_image = image[:, :, :3]
-        target_mask = None
-        if image.shape[2] == 4:
-            target_mask = image[:, :, 3].reshape((image.shape[0], image.shape[1], 1))
+        target_image = image
 
-        pad_lt = (64 * self.zoom) - (60 * self.zoom)
-        pad_rb = (64 * self.zoom) + (60 * self.zoom)
+        pad_lt = size // 32  # 8px on a 256px image
+        pad_rb = size - pad_lt
 
         warped_image = cv2.resize(  # pylint: disable=no-member
             warped_image[pad_lt:pad_rb, pad_lt:pad_rb, :],
-            (64, 64),
+            (64 * self.zoom[0], 64 * self.zoom[0]),
             cv2.INTER_AREA)  # pylint: disable=no-member
+        logger.trace("Warped image shape: %s", warped_image.shape)
         target_image = cv2.resize(  # pylint: disable=no-member
             target_image[pad_lt:pad_rb, pad_lt:pad_rb, :],
-            (64 * self.zoom, 64 * self.zoom),
+            (64 * self.zoom[1], 64 * self.zoom[1]),
             cv2.INTER_AREA)  # pylint: disable=no-member
-        if target_mask is None:
-            retval = warped_image, target_image
-        else:
+        logger.trace("Target image shape: %s", target_image.shape)
+
+        retval = [warped_image, target_image]
+
+        if mask is not None:
             target_mask = cv2.resize(  # pylint: disable=no-member
-                target_mask[pad_lt:pad_rb, pad_lt:pad_rb, :],
-                (64 * self.zoom, 64 * self.zoom),
+                mask[pad_lt:pad_rb, pad_lt:pad_rb, :],
+                (64 * self.zoom[1], 64 * self.zoom[1]),
                 cv2.INTER_AREA)  # pylint: disable=no-member
-            target_mask = target_mask.reshape((64 * self.zoom, 64 * self.zoom, 1))
-            retval = warped_image, target_image, target_mask
+            target_mask = target_mask.reshape((64 * self.zoom[1], 64 * self.zoom[1], 1))
+            logger.trace("Target mask shape: %s", target_mask.shape)
+
+            retval.append(target_mask)
+
         logger.trace("Randomly warped image")
         return retval
 

plugins/train/_config.py

@@ -82,6 +82,10 @@ def set_defaults(self):
                  "be named 'alignments.<file extension>' (eg. "
                  "alignments.json)."
                  "\nChoose from: 'json', 'pickle' or 'yaml'")
+        self.add_item(
+            section=section, title="mask_type", datatype=str, default="dfaker",
+            info="The mask to be used for training."
+                 "\nChoose from: 'dfaker' or 'dfl_full'")
 
         # << DFL MODEL OPTIONS >> #
         section = "dfl_h128"
@@ -92,6 +96,17 @@ def set_defaults(self):
             info="Lower memory mode. Set to 'True' if having issues with VRAM useage.\nNB: Models "
                  "with a changed lowmem mode are not compatible with each other."
                  "\nChoose from: True, False")
+        self.add_item(
+            section=section, title="alignments_format", datatype=str, default="json",
+            info="DFL-H128 model requires the alignments for your training "
+                 "images to be avalaible within the FACES folder.\nIt should "
+                 "be named 'alignments.<file extension>' (eg. "
+                 "alignments.json)."
+                 "\nChoose from: 'json', 'pickle' or 'yaml'")
+        self.add_item(
+            section=section, title="mask_type", datatype=str, default="dfl_full",
+            info="The mask to be used for training."
+                 "\nChoose from: 'dfaker' or 'dfl_full'")
 
         # << GAN MODEL OPTIONS >> #
         section = "gan_v2_2"

plugins/train/model/_base.py

@@ -9,6 +9,7 @@
 import sys
 
 from json import JSONDecodeError
+from keras import losses
 from keras.optimizers import Adam
 from keras.utils import multi_gpu_model
 
@@ -125,7 +126,7 @@ def loss_function(self):
         if self.config["dssim_loss"]:
             loss_func = DSSIMObjective()
         else:
-            loss_func = "mean_absolute_error"
+            loss_func = losses.mean_absolute_error
         logger.debug(loss_func)
         return loss_func
 

plugins/train/model/dfaker.py

@@ -26,10 +26,9 @@ def set_training_data(self):
         """ Set the dictionary for training """
         logger.debug("Setting training data")
         training_opts = dict()
-        serializer = self.config["alignments_format"]
-        training_opts["serializer"] = serializer
-        training_opts["use_mask"] = True
-        training_opts["use_alignments"] = True
+        training_opts["serializer"] = self.config["alignments_format"]
+        training_opts["mask_type"] = self.config["mask_type"]
+        training_opts["full_face"] = True
         training_opts["preview_images"] = 10
         logger.debug("Set training data: %s", training_opts)
         return training_opts

plugins/train/model/dfl_h128.py

@@ -11,7 +11,6 @@
 from .original import get_config, logger, Model as OriginalModel
 
 # TODO Implement DFL loss function (currently using dfaker)
-# TODO Get Mask working
 
 
 class Model(OriginalModel):
@@ -31,8 +30,7 @@ def set_training_data(self):
         """ Set the dictionary for training """
         logger.debug("Setting training data")
         training_opts = dict()
-        training_opts["use_mask"] = True
-        training_opts["remove_alpha"] = True
+        training_opts["mask_type"] = self.config["mask_type"]
         training_opts["preview_images"] = 10
         logger.debug("Set training data: %s", training_opts)
         return training_opts

plugins/train/model/original.py

@@ -46,7 +46,6 @@ def initialize(self):
             inp,
             self.networks["decoder_b"].network(self.networks["encoder"].network(inp)))
         self.add_predictors(ae_a, ae_b)
-
         logger.debug("Initialized model")
 
     def encoder(self):