From 9e45e2b6e9082c79d83c95759055e60333b7e41b Mon Sep 17 00:00:00 2001 From: MediaPipe Team Date: Wed, 16 Aug 2023 09:59:53 -0700 Subject: [PATCH] Setting training for the encoder and decoder when converting to TFLite. Also add selected TF ops to TFLite converter. PiperOrigin-RevId: 557520277 --- .../python/vision/face_stylizer/face_stylizer.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/mediapipe/model_maker/python/vision/face_stylizer/face_stylizer.py b/mediapipe/model_maker/python/vision/face_stylizer/face_stylizer.py index 09b5ab8c0..147cef004 100644 --- a/mediapipe/model_maker/python/vision/face_stylizer/face_stylizer.py +++ b/mediapipe/model_maker/python/vision/face_stylizer/face_stylizer.py @@ -231,8 +231,8 @@ class FaceStylizer(object): # Create an end-to-end model by concatenating encoder and decoder inputs = tf.keras.Input(shape=(256, 256, 3)) - x = self._encoder(inputs) - x = self._decoder({'inputs': x + self.w_avg}) + x = self._encoder(inputs, training=True) + x = self._decoder({'inputs': x + self.w_avg}, training=True) x = x['image'][-1] # Scale the data range from [-1, 1] to [0, 1] to support running inference # on both CPU and GPU. @@ -241,6 +241,10 @@ class FaceStylizer(object): face_stylizer_model_buffer = model_util.convert_to_tflite( model=model, + supported_ops=[ + tf.lite.OpsSet.TFLITE_BUILTINS, + tf.lite.OpsSet.SELECT_TF_OPS, + ], preprocess=self._preprocessor, )