Model output not looking as expected (tfjs/tflite) styletransfer model

General Discussion
, ,
#1

Hey all, I’m trying to implement the style transfer from the TFHub: TensorFlow Hub using tf.js but I’m running into a little bit of an issue. I think I’ve got it all hooked up together but the resulting image is not looking very good. Any help would be truly appreciated!

Current result in the browser

652×1042 924 KB

Code

import '@tensorflow/tfjs-backend-cpu';
import{ browser, image, tidy, Tensor, reshape} from '@tensorflow/tfjs-core';
import * as tflite from '@tensorflow/tfjs-tflite'

const predictionUrl = 'https://tfhub.dev/google/lite-model/magenta/arbitrary-image-stylization-v1-256/fp16/prediction/1';
const transferUrl = 'https://tfhub.dev/google/lite-model/magenta/arbitrary-image-stylization-v1-256/fp16/transfer/1';

export async function run() {
  const predictionModel = await tflite.loadTFLiteModel(predictionUrl);
  const transferMode = await tflite.loadTFLiteModel(transferUrl);
  const styleImg = await browser.fromPixelsAsync(document.querySelector('.styleImage'));
  const img = await browser.fromPixelsAsync(document.querySelector('.originalImage'));
  const resizedStyleImage = reshape(image.resizeBilinear(styleImg, [256, 256]), [1, 256, 256, 3]);
  const styleBottleNeck = await tidy(() => predictionModel.predict([resizedStyleImage])) as Tensor;
  const resizedImageTensor =  reshape(image.resizeBilinear(img, [384, 384], true), [1, 384, 384, 3]);
  const output = await tidy(() => transferMode.predict([resizedImageTensor, styleBottleNeck])) as Tensor;
  return browser.toPixels(reshape(image.resizeBilinear(output, [540, 960], true),  [540, 960, 3]), document.querySelector('.canvas'))
}
1 Like
#3

Hi, can you provide the result from TF using their colab?

#4

This is what I get using the CoLab:

index
index1178×410 686 KB

#5

Is there maybe some preprocessing step I’m missing in my code?

#6

In the hub example I see:

Convert to float32 numpy array, add batch dimension, and normalize to range [0, 1]
1 Like
#7

Yes, you need to normalize the input image tensor to [0, 1].

1 Like
#9

That seems to have vastly improved the result. I still need to validate that the results match with the CoLab given the exact same images as I may be missing additional processing steps.

For folks who may run into similar issues, to normalize, I divided the image tensors by 255 like so:

const normalizedStyleImg= div(await browser.fromPixelsAsync(document.querySelector('.styleImage')), 255);
const normalizedContentImg = div(await browser.fromPixelsAsync(document.querySelector('.originalImage')), 255);
3 Likes
#11

Do let us know if you end up putting the demo somewhere! Would love to share with the community if it is public :slight_smile: Just tag on social with #MadeWithTFJS so we can find it if it is!

2 Likes