Add a face alignment preprocessor to face stylizer.

PiperOrigin-RevId: 542559764

mediapipe/model_maker/python/vision/face_stylizer/BUILD

@@ -20,13 +20,6 @@ licenses(["notice"])
 
 package(default_visibility = ["//mediapipe:__subpackages__"])
 
-filegroup(
-    name = "testdata",
-    srcs = glob([
-        "testdata/**",
-    ]),
-)
-
 py_library(
     name = "constants",
     srcs = ["constants.py"],
@@ -72,18 +65,11 @@ py_library(
     name = "dataset",
     srcs = ["dataset.py"],
     deps = [
+        ":constants",
         "//mediapipe/model_maker/python/core/data:classification_dataset",
-        "//mediapipe/model_maker/python/vision/core:image_utils",
-    ],
-)
-
-py_test(
-    name = "dataset_test",
-    srcs = ["dataset_test.py"],
-    data = [":testdata"],
-    deps = [
-        ":dataset",
-        "//mediapipe/tasks/python/test:test_utils",
+        "//mediapipe/python:_framework_bindings",
+        "//mediapipe/tasks/python/core:base_options",
+        "//mediapipe/tasks/python/vision:face_aligner",
     ],
 )
 

mediapipe/model_maker/python/vision/face_stylizer/constants.py

@@ -41,5 +41,11 @@ FACE_STYLIZER_W_FILES = file_util.DownloadedFiles(
     'https://storage.googleapis.com/mediapipe-assets/face_stylizer_w_avg.npy',
 )
 
+FACE_ALIGNER_TASK_FILES = file_util.DownloadedFiles(
+    'face_stylizer/face_landmarker_v2.task',
+    'https://storage.googleapis.com/mediapipe-assets/face_landmarker_v2.task',
+    is_folder=False,
+)
+
 # Dimension of the input style vector to the decoder
 STYLE_DIM = 512

mediapipe/model_maker/python/vision/face_stylizer/dataset.py

@@ -13,13 +13,37 @@
 # limitations under the License.
 """Face stylizer dataset library."""
 
+from typing import Sequence
 import logging
 import os
 
 import tensorflow as tf
 
 from mediapipe.model_maker.python.core.data import classification_dataset
-from mediapipe.model_maker.python.vision.core import image_utils
+from mediapipe.model_maker.python.vision.face_stylizer import constants
+from mediapipe.python._framework_bindings import image as image_module
+from mediapipe.tasks.python.core import base_options as base_options_module
+from mediapipe.tasks.python.vision import face_aligner
+
+
+def _preprocess_face_dataset(
+    all_image_paths: Sequence[str],
+) -> Sequence[tf.Tensor]:
+  """Preprocess face image dataset by aligning the face."""
+  path = constants.FACE_ALIGNER_TASK_FILES.get_path()
+  base_options = base_options_module.BaseOptions(model_asset_path=path)
+  options = face_aligner.FaceAlignerOptions(base_options=base_options)
+  aligner = face_aligner.FaceAligner.create_from_options(options)
+
+  preprocessed_images = []
+  for path in all_image_paths:
+    tf.compat.v1.logging.info('Preprocess image %s', path)
+    image = image_module.Image.create_from_file(path)
+    aligned_image = aligner.align(image)
+    aligned_image_tensor = tf.convert_to_tensor(aligned_image.numpy_view())
+    preprocessed_images.append(aligned_image_tensor)
+
+  return preprocessed_images
 
 
 # TODO: Change to a unlabeled dataset if it makes sense.
@@ -58,6 +82,7 @@ class Dataset(classification_dataset.ClassificationDataset):
     ):
       raise ValueError('No images found under given directory')
 
+    image_data = _preprocess_face_dataset(all_image_paths)
     label_names = sorted(
         name
         for name in os.listdir(data_root)
@@ -73,11 +98,7 @@ class Dataset(classification_dataset.ClassificationDataset):
         for path in all_image_paths
     ]
 
-    path_ds = tf.data.Dataset.from_tensor_slices(all_image_paths)
-
-    image_ds = path_ds.map(
-        image_utils.load_image, num_parallel_calls=tf.data.AUTOTUNE
-    )
+    image_ds = tf.data.Dataset.from_tensor_slices(image_data)
 
     # Load label
     label_ds = tf.data.Dataset.from_tensor_slices(

mediapipe/model_maker/python/vision/face_stylizer/dataset_test.py

@@ -12,8 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import numpy as np
 import tensorflow as tf
 
+from mediapipe.model_maker.python.vision.core import image_utils
 from mediapipe.model_maker.python.vision.face_stylizer import dataset
 from mediapipe.tasks.python.test import test_utils
 
@@ -22,10 +24,10 @@ class DatasetTest(tf.test.TestCase):
 
   def setUp(self):
     super().setUp()
-    self._test_data_dirname = 'input/style'
 
   def test_from_folder(self):
-    input_data_dir = test_utils.get_test_data_path(self._test_data_dirname)
+    test_data_dirname = 'input/style'
+    input_data_dir = test_utils.get_test_data_path(test_data_dirname)
     data = dataset.Dataset.from_folder(dirname=input_data_dir)
     self.assertEqual(data.num_classes, 2)
     self.assertEqual(data.label_names, ['cartoon', 'sketch'])

mediapipe/model_maker/python/vision/face_stylizer/face_stylizer.py

@@ -14,7 +14,7 @@
 """APIs to train face stylization model."""
 
 import os
-from typing import Callable, Optional
+from typing import Any, Callable, Optional
 
 import numpy as np
 import tensorflow as tf
@@ -54,7 +54,6 @@ class FaceStylizer(object):
     self._model_spec = model_spec
     self._model_options = model_options
     self._hparams = hparams
-    # TODO: Support face alignment in image preprocessor.
     self._preprocessor = image_preprocessing.Preprocessor(
         input_shape=self._model_spec.input_image_shape,
         num_classes=1,
@@ -128,7 +127,7 @@ class FaceStylizer(object):
   def _train_model(
       self,
       train_data: classification_ds.ClassificationDataset,
-      preprocessor: Optional[Callable[..., bool]] = None,
+      preprocessor: Optional[Callable[..., Any]] = None,
   ):
     """Trains the face stylizer model.