Hey all, I’m trying to implement the style transfer from the TFHub: Google | arbitrary-image-stylization-v1 | Kaggle 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
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'))
}