My TFJS stylization looks terrible

[Ryan_Johnson]

I tried to convert this guide into TFJS: Google Colab. I took this model: Google | arbitrary-image-stylization-v1 | Kaggle and used the tensorflowjs_converter to convert it to a tfjs_graph_model using this command:

tensorflowjs_converter --signature=serving_default --input_format=tf_hub --output_format=tfjs_graph_model ~/Downloads/magenta_arbitrary-image-stylization-v1-256_2 ../temp/test_model

I then put that folder on my test server, along with the example images of the golden gate bridge and the wave. I wrote a small amount of js and the result was terrible. Here is what I wrote:

const styleImg = tf.div(await tf.browser.fromPixelsAsync(document.querySelector('.styleImage')), 255); 
const img = tf.div(await tf.browser.fromPixelsAsync(document.querySelector('.originalImage')), 255);
const resizedStyleImage = tf.reshape(tf.image.resizeBilinear(styleImg, [256, 256]), [1, 256, 256, 3]);
const resizedImageTensor =  tf.reshape(tf.image.resizeBilinear(img, [288, 384], true), [1, 288, 384, 3]);
var mymodel = await tf.loadGraphModel(TEST_MODEL_URL)
test = await mymodel.execute([resizedImageTensor, resizedStyleImage]);
tf.browser.toPixels(tf.reshape(tf.image.resizeBilinear(test, [288, 384], true), [288, 384, 3]), document.querySelector('.my-canvas'));

Any thoughts on where I went wrong? Perhaps this is a side effect of converting the model from hub to tfjs_graph? I’m new to ML so it’s possible I made an obvious mistake and I’m just oblivious to it.

[Ryan_Johnson]

I figured it out! My mistake was I swapped the inputs around in my execute function. So I was making a wave in the style of a bridge, not a bridge in the style of a wave. Here is my final result, which I am much happier with:

Corrected line of code I ended up using:

test = await mymodel.predict([resizedStyleImage, resizedImageTensor])

After fixing the issue, I did see this warning in my console, so perhaps I’m still doing something wrong?

High memory usage in GPU: 1211.99 MB, most likely due to a memory leak

I think I need to wrap predict in a tf.tidy perhaps.