How to convert an Image tensor to PNG using react-native?

Question

I am using react-native-tfjs for running a neural network, it gives me an RGB image tensor as one of the outputs. How can I save this as a PNG image?

What have I tried so far:

• Tried to use the tf.browser.toPixels function, but this does not work as for drawing or saving it requires html5 canvas element.
• Searched for EncodePng implementations, tf.node.encodePng unfortunately there are no similar implementations for react-native.
• Bunch of other things like trying to convert the buffer to base64 but all with no luck.

Here is the implementation of Jpeg encoding that I found that works:

import * as tf from "@tensorflow/tfjs";
import * as jpeg from "jpeg-js";

async function tensorToImageUrl(imageTensor) {
  const [height, width] = imageTensor.shape;
  const buffer = await imageTensor.toInt().data();
  const frameData = new Uint8Array(width * height * 4);

  let offset = 0;
  for (let i = 0; i < frameData.length; i += 4) {
    frameData[i] = buffer[offset];
    frameData[i + 1] = buffer[offset + 1];
    frameData[i + 2] = buffer[offset + 2];
    frameData[i + 3] = 0xff;

    offset += 3;
  }

  const rawImageData = {
    data: frameData,
    width,
    height,
  };
  const jpegImageData = jpeg.encode(rawImageData, 100);
  const base64Encoding = tf.util.decodeString(jpegImageData.data, "base64");
  return base64Encoding;
}

I have to have a lossless image for further retrieval & processing, so I cannot use jpeg-js, I tried to look for similar libraries for png but a lot of them are outdated or lack documentation, if anyone could help me implemetnt the tensorToImageUrl function but for PNGs, it would be very helpful for my project.

Answer 1

Ok so I managed to get it work, hopefully it helps someone in the future and other similar unanswered questions. I used Jimp Image processing library for it, it is pure written in pure JS and works in Node and browser but there is a hack to get it to work in react-native/react.

import * as tf from "@tensorflow/tfjs";

import * as _Jimp from "jimp";
const Jimp = typeof self !== "undefined" ? self.Jimp || _Jimp : _Jimp; // The hack

async function encodePng(imageTensor) {
  const [height, width] = imageTensor.shape;
  const buffer = await imageTensor.toInt().data();
  const frameData = new Uint8Array(width * height * 4);

  let offset = 0;
  for (let i = 0; i < frameData.length; i += 4) {
    frameData[i] = buffer[offset];
    frameData[i + 1] = buffer[offset + 1];
    frameData[i + 2] = buffer[offset + 2];
    frameData[i + 3] = 0xff; // Buffer has to be Uint8Array in RGBA raw data because Jimp expects the format, we stuff dummy alpha values here.

    offset += 3;
  }
  const rawImageData = {
    data: frameData,
    width,
    height,
  };
  var base64Data = "";
  try {
    const image = await Jimp.read(rawImageData);
    console.log(image);
    base64Data = await image.getBase64Async(Jimp.MIME_PNG); // JIMP support Jpeg, BMP, PNG etc. so you can modify this however you like.
  } catch (e) {
    console.log(e);
  }
  if (base64Data === "") console.log("operation unsucessful");
  return base64Data;
}