local tfjs models assets
npm install local-tfjs-modelsThis is a collection repo of tfjs-models. Then you can deploy these models in your local or cdn enviroment.
typescript
import * as tf from '@tensorflow/tfjs';const STYLE_MODEL_URL_MAP = {
"hayao": "https://unpkg.com/local-tfjs-models@0.0.3/cartoon-GAN/hayao/model.json",
"hosoda": "https://unpkg.com/local-tfjs-models@0.0.3/cartoon-GAN/hosoda/model.json",
"paprika": "https://unpkg.com/local-tfjs-models@0.0.3/cartoon-GAN/paprika/model.json",
"shinkai": "https://unpkg.com/local-tfjs-models@0.0.3/cartoon-GAN/shinkai/model.json"
};
type STYLE_TYPE = "hayao" | "hosoda" | "paprika" | "shinkai";
async function setupModel(style: STYLE_TYPE): Promise {
const model = await tf.loadGraphModel(STYLE_MODEL_URL_MAP[style]);
// this predict action is used to save time, because every first predict action is very slow.
model.predict(tf.zeros([1, 1, 1, 3]));
return model;
}
function predict(
model: tf.GraphModel,
inputImgElement: HTMLImageElement | HTMLCanvasElement,
outputElement: HTMLCanvasElement
) {
// convert the input image to tensor accepted by model
let inputTensor = tf.browser.fromPixels(inputImgElement);
inputTensor = inputTensor.toFloat();
inputTensor = inputTensor.reverse(2);
inputTensor = tf.expandDims(inputTensor, 0);
let outputTensor = model.predict(inputTensor as tf.Tensor) as tf.Tensor;
// convert the predict tensor to output image
outputTensor = tf.squeeze(outputTensor, [0]);
outputTensor = outputTensor.reverse(2);
outputTensor = outputTensor.mul(0.5).add(0.5);
outputTensor = tf.clipByValue(outputTensor, 0, 1);
// put the result into canvas element.
tf.browser.toPixels(outputTensor as tf.Tensor2D, outputElement);
}
// main
/*
* style: the name of style model
* img: the input image/canvas element
* out: the output canvas element
*/
// load and init model
const model = await setupModel(style = "hayao");
// predict the result and paint it to the output canvas element
predict(model, img, out);
``