mediapipe_model_maker.face_stylizer.FaceStylizer

FaceStylizer for building face stylization model.

mediapipe_model_maker.face_stylizer.FaceStylizer(
    model_spec: mediapipe_model_maker.face_stylizer.model_spec.ModelSpec,
    model_options: mediapipe_model_maker.face_stylizer.ModelOptions,
    hparams: mediapipe_model_maker.face_stylizer.HParams
)

model_spec Specification for the model.
model_options Model options for creating face stylizer.
hparams The hyperparameters for training face stylizer.

w_avg An average face latent code to regularize face generation in face stylization.

Methods

create

View source

@classmethod
create(
    train_data: mediapipe_model_maker.face_stylizer.dataset.classification_dataset.ClassificationDataset,
    options: mediapipe_model_maker.face_stylizer.FaceStylizerOptions
) -> 'FaceStylizer'

Creates and trains a face stylizer with input datasets.

Args
train_data The input style image dataset for training the face stylizer.
options The options to configure face stylizer.

Returns
A FaceStylizer instant with the trained model.

export_model

View source

export_model(
    model_name: str = 'face_stylizer.task'
)

Converts the model to TFLite and exports as a model bundle file.

Saves a model bundle file and metadata json file to hparams.export_dir. The resulting model bundle file will contain necessary models for face detection, face landmarks detection, and customized face stylization. Only the model bundle file is needed for the downstream face stylization task. The metadata.json file is saved only to interpret the contents of the model bundle file. The face detection model and face landmarks detection model are from https://storage.googleapis.com/mediapipe-assets/face_landmarker_v2.task and the customized face stylization model is trained in this library.

Args
model_name Face stylizer model bundle file name. The full export path is {self._hparams.export_dir}/{model_name}.

stylize

View source

stylize(
    data: mediapipe_model_maker.face_stylizer.dataset.classification_dataset.ClassificationDataset
) -> mediapipe_model_maker.face_stylizer.dataset.classification_dataset.ClassificationDataset

Stylizes the images represented by the input dataset.

Args
data Dataset of input images, can contain multiple images.

Returns
A dataset contains the stylized images