How to Create a Cartoonizer with TensorFlow Lite?

Create the TensorFlow Lite Model

Generate a SavedModel from the pre-trained model weights

├── checkpoint├── model-33999.data-00000-of-00001└── model-33999.index

Convert SavedModel to TFLite

# Create a concrete function from the SavedModel
model = tf.saved_model.load(saved_model_dir)
concrete_func = model.signatures[
tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY]
# Specify the input shape
concrete_func.inputs[0].set_shape([1, IMG_SHAPE, IMG_SHAPE, 3])
# Convert the model and export
converter = tf.lite.TFLiteConverter.from_concrete_functions([concrete_func])
converter.optimizations = [tf.lite.Optimize.DEFAULT]
converter.target_spec.supported_types = [tf.float16] # Only for float16
tflite_model = converter.convert()
open(tflite_model_path, ‘wb’).write(tflite_model)

Running inference in Python

interpreter = tf.lite.Interpreter(model_path=tflite_model_path)
input_details = interpreter.get_input_details()
interpreter.allocate_tensors()
interpreter.set_tensor(input_details[0][‘index’], preprocessed_source_image)
interpreter.invoke()
raw_prediction = interpreter.tensor(
interpreter.get_output_details()[0][‘index’])()

Add metadata for easy integration with a mobile app

input_image_normalization.options.mean = [127.5]
input_image_normalization.options.std = [127.5]
output_image_normalization.options.mean = [-1]
output_image_normalization.options.std = [0.00784313] # 1/127.5

Benchmark models on Android (Optional)

bazel build -c opt \
— config=android_arm64 \
tensorflow/lite/tools/benchmark:benchmark_model
adb push benchmark_model /data/local tmp
adb shell chmod +x /data/local/tmp/benchmark_model
adb push whitebox_cartoon_gan_dr.tflite /data/local/tmp
adb shell /data/local/tmp/android_aarch64_benchmark_model \
— graph=/data/local/tmp/whitebox_cartoon_gan_dr.tflite \
— num_threads=4

Model deployment to Android

Download Android Studio 4.1 Preview

Create a new Android Project

Set up CameraX for image capture

Import the TensorFlow Lite models

Putting everything together

Cartoonizer Android app

Acknowledgments

Google Developer Expert for ML | AI, Art & Design | margaretmz.art

Get the Medium app

A button that says 'Download on the App Store', and if clicked it will lead you to the iOS App store
A button that says 'Get it on, Google Play', and if clicked it will lead you to the Google Play store