从张量到jpeg的tfjs-react-native

时间:2020-10-01 07:30:00

标签: react-native jpeg tensorflow.js jpeg-js

decodeJpeg中有@tensorflow/tfjs-react-native,但没有encodeJpeg。然后如何将张量写入本地jpeg文件中?

我试图看一下代码并“反转”函数,最后我写了:

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
}

但是当我显示带有<Image />的图像时,我看到所有的都是纯绿色。

2 个答案:

答案 0 :(得分:0)

您可以将imgBase64直接用于图像组件,如下所示:

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

答案 1 :(得分:0)

这是我这样做的最终工具:

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}
}