mediapipe_model_maker.image_classifier.ImageClassifier

ImageClassifier for building image classification model.

mediapipe_model_maker.image_classifier.ImageClassifier( model_spec: mediapipe_model_maker.image_classifier.ModelSpec, label_names: List[str], hparams: mediapipe_model_maker.image_classifier.HParams, model_options: mediapipe_model_maker.image_classifier.ModelOptions )

model_spec Specification for the model.
label_names A list of label names for the classes.
hparams The hyperparameters for training image classifier.
model_options Model options for creating image classifier.

Methods

create

View source

@classmethod create( train_data: mediapipe_model_maker.face_stylizer.dataset.classification_dataset.ClassificationDataset, validation_data: mediapipe_model_maker.face_stylizer.dataset.classification_dataset.ClassificationDataset, options: mediapipe_model_maker.image_classifier.ImageClassifierOptions ) -> 'ImageClassifier'

Creates and trains an ImageClassifier.

Loads data and trains the model based on data for image classification. If a checkpoint file exists in the {options.hparams.export_dir}/checkpoint/ directory, the training process will load the weight from the checkpoint file for continual training.

Args
train_data Training data.
validation_data Validation data.
options configuration to create image classifier.

Returns
An instance based on ImageClassifier.

evaluate

View source

evaluate( data: mediapipe_model_maker.model_util.dataset.Dataset, batch_size: int = 32 ) -> Any

Evaluates the classifier with the provided evaluation dataset.

Args
data Evaluation dataset
batch_size Number of samples per evaluation step.

Returns
The loss value and accuracy.

export_labels

View source

export_labels( export_dir: str, label_filename: str = 'labels.txt' )

Exports classification labels into a label file.

Args
export_dir The directory to save exported files.
label_filename File name to save labels model. The full export path is {export_dir}/{label_filename}.

export_model

View source

export_model( model_name: str = 'model.tflite', quantization_config: Optional[mediapipe_model_maker.quantization.QuantizationConfig] = None )

Converts and saves the model to a TFLite file with metadata included.

Note that only the TFLite file is needed for deployment. This function also saves a metadata.json file to the same directory as the TFLite file which can be used to interpret the metadata content in the TFLite file.

Args
model_name File name to save TFLite model with metadata. The full export path is {self._hparams.export_dir}/{model_name}.
quantization_config The configuration for model quantization.

export_tflite

View source

export_tflite( export_dir: str, tflite_filename: str = 'model.tflite', quantization_config: Optional[mediapipe_model_maker.quantization.QuantizationConfig] = None, preprocess: Optional[Callable[..., bool]] = None )

Converts the model to requested formats.

Args
export_dir The directory to save exported files.
tflite_filename File name to save TFLite model. The full export path is {export_dir}/{tflite_filename}.
quantization_config The configuration for model quantization.
preprocess A callable to preprocess the representative dataset for quantization. The callable takes three arguments in order: feature, label, and is_training.

summary

View source

summary()

Prints a summary of the model.