tfjs-react-native from tensor to jpeg

Question

There is decodeJpeg from @tensorflow/tfjs-react-native but no encodeJpeg. How to write a tensor into a local jpeg file then?

I have tried to look at the code and "inverse" the function, I ended up writing:

import * as tf from '@tensorflow/tfjs';
import * as FileSystem from 'expo-file-system';
import * as jpeg from 'jpeg-js';

export const encoderJpeg = async (tensor, name) => {
  // add alpha channel if missing
  const shape = [...tensor.shape]
  shape.pop()
  shape.push(4)
  const tensorWithAlpha = tf.concat([tensor, tensor], [-1]).slice([0], shape)
  
  const array = new Uint8Array(tensorWithAlpha.dataSync())
  const rawImageData = {
    data: array.buffer,
    width: shape[1],
    height: shape[0],
  };
  const jpegImageData = jpeg.encode(rawImageData, 50);
  const imgBase64 = tf.util.decodeString(jpegImageData.data, "base64")
  const uri = FileSystem.documentDirectory + name;
  await FileSystem.writeAsStringAsync(uri, imgBase64, {
    encoding: FileSystem.EncodingType.Base64,
  });
  return uri
}

but when I show the images with an <Image /> I see that there are all plain green.

Answer 1

You can use imgBase64 directly into your image component as follows:

<Image source={{uri: 'data:image/jpeg;base64,' + imgBase64}} />

Answer 2

This is my final util for doing this:

import * as tf from '@tensorflow/tfjs';
import * as FileSystem from 'expo-file-system';
import * as jpeg from 'jpeg-js';

export const encodeJpeg = async (tensor) => {

  const height = tensor.shape[0]
  const width = tensor.shape[1]
  const data = new Buffer(
    // concat with an extra alpha channel and slice up to 4 channels to handle 3 and 4 channels tensors
    tf.concat([tensor, tf.ones([height, width, 1]).mul(255)], [-1])
      .slice([0], [height, width, 4])
      .dataSync(),
  )

  const rawImageData = {data, width, height};
  const jpegImageData = jpeg.encode(rawImageData, 100);

  const imgBase64 = tf.util.decodeString(jpegImageData.data, "base64")
  const salt = `${Date.now()}-${Math.floor(Math.random() * 10000)}`
  const uri = FileSystem.documentDirectory + `tensor-${salt}.jpg`;
  await FileSystem.writeAsStringAsync(uri, imgBase64, {
    encoding: FileSystem.EncodingType.Base64,
  });
  return {uri, width, height}
}