Hey all, I’m trying to implement the style transfer from the TFHub: TensorFlow Hub 2 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'))
}