Convert Pix2Pix models to TensorFlow.js

Here's how to convert Pix2Pix models trained with TensorFlow to TensorFlow.js for use on the web or with Node.js with @tensorflow/tfjs-node and/or @tensorflow/tfjs-core.

Training a Pix2Pix model with TensorFlow

The first thing we need is to train a Pix2Pix model with TensorFlow. You can watch my video on how to Train and Predict with Pix2Pix.

That workflow will produce an .h5 file¹.

Converting a Pix2Pix model to TensorFlow.js format

The .h5 models can be loaded with TensorFlow Core but won't be usable in the JavaScript version of TensorFlow.

For conversion, TensorFlow provides the tfjs-converter library as a "a command line utility that converts Keras and TensorFlow models for use in TensorFlow.js."

Here's a sample command you can use to convert a model.

tensorflowjs_converter \
    --input_format=keras \
    /tmp/model.h5 \
    /tmp/tfjs_model

Installing the tfjs-converter library

Optionally, we can create an Anaconda environment to install these dependencies we can discard after we're done converting our models. Note that, as of October 11, 2022, the tfjs-converter repository specifies Python 3.6.8 as a requirement for the converter to work, but I was able to make it work with Python 3.9.

# Create a new environment (optional)
conda create -n tfjs python=3.9 -y

# Activate conda environment (optional)
conda activate tfjs

# Install tfjs-converter
pip install tensorflowjs

Now we're all set to convert our H5 model to TensorFlow.js.

tensorflowjs_converter \
    --input_format=keras \
    ./model.h5 \
    ./tfjs-model

If you created an Anaconda environment, here's how to dispose of it. You can also keep it for future model conversions.

# Exit the environment
conda deactivate

# Delete conda environment
conda remove -n tfjs --all -y

Cannot install tensorflowjs

You may get the following error in Apple Silicon Macs because the tensorflow dependency package can't be resolved. I created a Google Colab notebook so you can test this on the browser.

ERROR: Cannot install tensorflowjs==0.1.0, tensorflowjs==0.1.1, tensorflowjs==0.1.2, tensorflowjs==0.2.0, tensorflowjs==0.2.1, tensorflowjs==0.3.0, tensorflowjs==0.3.1, tensorflowjs==0.4.0, tensorflowjs==0.4.1, tensorflowjs==0.4.2, tensorflowjs==0.5.0, tensorflowjs==0.5.2, tensorflowjs==0.5.4, tensorflowjs==0.5.6, tensorflowjs==0.5.7, tensorflowjs==0.6.0, tensorflowjs==0.6.1, tensorflowjs==0.6.2, tensorflowjs==0.6.4, tensorflowjs==0.6.5, tensorflowjs==0.6.7, tensorflowjs==0.8.0, tensorflowjs==0.8.5, tensorflowjs==0.8.6, tensorflowjs==1.0.1, tensorflowjs==1.1.2, tensorflowjs==1.2.1, tensorflowjs==1.2.10, tensorflowjs==1.2.10.1, tensorflowjs==1.2.2, tensorflowjs==1.2.2.1, tensorflowjs==1.2.3, tensorflowjs==1.2.6, tensorflowjs==1.2.9, tensorflowjs==1.3.1, tensorflowjs==1.3.1.1, tensorflowjs==1.3.2, tensorflowjs==1.4.0, tensorflowjs==1.5.2, tensorflowjs==1.6.0, tensorflowjs==1.7.2, tensorflowjs==1.7.3, tensorflowjs==1.7.4, tensorflowjs==1.7.4.post1, tensorflowjs==2.0.0, tensorflowjs==2.0.1, tensorflowjs==2.0.1.post1, tensorflowjs==2.1.0, tensorflowjs==2.3.0, tensorflowjs==2.4.0, tensorflowjs==2.5.0, tensorflowjs==2.6.0, tensorflowjs==2.7.0, tensorflowjs==2.8.0, tensorflowjs==2.8.1, tensorflowjs==2.8.2, tensorflowjs==2.8.3, tensorflowjs==2.8.4, tensorflowjs==2.8.5, tensorflowjs==3.0.0, tensorflowjs==3.1.0, tensorflowjs==3.11.0, tensorflowjs==3.12.0, tensorflowjs==3.13.0, tensorflowjs==3.14.0, tensorflowjs==3.15.0, tensorflowjs==3.17.0, tensorflowjs==3.18.0, tensorflowjs==3.19.0, tensorflowjs==3.2.0, tensorflowjs==3.20.0, tensorflowjs==3.21.0, tensorflowjs==3.3.0, tensorflowjs==3.4.0, tensorflowjs==3.5.0, tensorflowjs==3.6.0, tensorflowjs==3.7.0, tensorflowjs==3.8.0 and tensorflowjs==3.9.0 because these package versions have conflicting dependencies.