如何在React Native中将图像张量转换为PNG格式？

我正在使用react-native-tfjs运行神经网络，其输出之一为RGB图像张量，需要将其保存为PNG图像，但尝试多种方法均未成功。

已尝试的方案 #

• 调用tf.browser.toPixels：该函数依赖HTML5 Canvas元素，无法在React Native环境中生效
• 查找tf.node.encodePng：React Native中没有对应的实现
• 尝试将张量缓冲区转换为base64：未成功实现

可用的JPEG编码实现（参考） #

以下是可正常运行的JPEG编码逻辑，由于需求为无损图像，无法采用该方案：

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

PNG编码实现方案 #

可以使用纯JavaScript的PNG编码库upng-js来实现，无需原生依赖，适配React Native环境：

步骤1：安装依赖 #

npm install upng-js
# 或
yarn add upng-js

步骤2：实现PNG转换函数 #

import * as tf from "@tensorflow/tfjs";
import UPNG from 'upng-js';

async function tensorToPngBase64(imageTensor) {
  // 确保张量值在0-255范围内，并转换为整数类型
  const normalizedTensor = tf.clipByValue(imageTensor, 0, 255).toInt();
  const [height, width] = normalizedTensor.shape;

  // 获取RGB格式的张量数据
  const rgbData = await normalizedTensor.data();
  // 构造RGBA格式数据（添加不透明的Alpha通道）
  const rgbaData = new Uint8Array(width * height * 4);

  let rgbOffset = 0;
  for (let i = 0; i < rgbaData.length; i += 4) {
    rgbaData[i] = rgbData[rgbOffset];     // 红色通道
    rgbaData[i + 1] = rgbData[rgbOffset + 1]; // 绿色通道
    rgbaData[i + 2] = rgbData[rgbOffset + 2]; // 蓝色通道
    rgbaData[i + 3] = 0xff;               // 不透明Alpha通道
    rgbOffset += 3;
  }

  // 编码为PNG二进制数据
  const pngBuffer = UPNG.encode([rgbaData], width, height, 0);
  // 转换为base64编码字符串
  return tf.util.encodeString(pngBuffer, 'base64');
}

使用说明 #

• 确保输入张量的shape为[height, width, 3]（RGB格式），若包含batch维度（如[batchSize, height, width, 3]），需先提取单个样本：imageTensor = imageTensor[0]
• 转换得到的base64字符串可直接用于图像显示（格式为data:image/png;base64,${base64Result}），或通过react-native-fs等库保存为本地PNG文件

内容的提问来源于stack exchange，提问作者Sameer Trivedi