tfjs-models

Some tfjs models for frontend ml dev purpose.

progan

Progressive GAN trained on CelebA for 128x128 images.

model

https://tfhub.dev/google/progan-128/1

convert

tensorflowjs_converter \
    --input_format=tf_hub \
    'https://tfhub.dev/google/progan-128/1' \
    ./progan

usage

import * as tf from "@tensorflow/tfjs-node";
import { Tensor3D } from "@tensorflow/tfjs-node";

const fs = require('fs');
const handler = tf.io.fileSystem("./model/progan/model.json");

async function main() {
  const model = await tf.loadGraphModel(handler);
  const res = model.predict(tf.randomNormal([1, 512])) as tf.Tensor;

  let outputTensor = tf.squeeze(res, [0]) as tf.Tensor<tf.Rank>;

  outputTensor = outputTensor.mul(255);
  outputTensor = tf.clipByValue(outputTensor, 0, 255);

  const outputImage = await tf.node.encodeJpeg(outputTensor as Tensor3D);
  fs.writeFileSync('gan.jpeg', outputImage); 
}

main();

esrgan

Enhanced Super Resolution GAN (Wang et. al.)[1] for image super resolution. Produces x4 Super Resolution Image from images of {Height, Width} >=64. Works best on Bicubically downsampled images.\ (This is because, the model is originally trained on Bicubically Downsampled DIV2K Dataset)

model

https://tfhub.dev/captain-pool/esrgan-tf2/1

convert

mkdir ersgan_saved_model ersgan
cd ersgan_saved_model
wget https://storage.googleapis.com/tfhub-modules/captain-pool/esrgan-tf2/1.tar.gz
tar zxvf 1.tar.gz
cd ../

tensorflowjs_converter \
    --input_format=tf_saved_model \
    /ersgan_saved_model/saved_model.pb \
    /ersgan

usage

import * as tf from "@tensorflow/tfjs-node";
import { Tensor3D } from "@tensorflow/tfjs-node";

const fs = require('fs');
const handler = tf.io.fileSystem("./model/esrgan/model.json");

async function main() {
  const model = await tf.loadGraphModel(handler);

  let img = fs.readFileSync(process.argv[2] || "chin.png");
  const im = tf.node.decodeJpeg(img).toFloat().expandDims();
  const res = model.predict(im) as tf.Tensor;;

  let outputTensor = tf.squeeze(res, [0]) as tf.Tensor<tf.Rank>;

  const outputImage = await tf.node.encodeJpeg(outputTensor as Tensor3D);
  fs.writeFileSync('esrgan.jpeg', outputImage); 
}

main();