Flutter中Tflite模型输出TensorBuffer转TensorImage/Image问题

我用Flutter的tflite库做图像分割，输入是RGB格式[3,224,224]的张量，模型运行正常，但输出的[1,1,224,224]格式TensorBuffer没法转成TensorImage或普通Image：

• 调用TensorImage.fromTensorBuffer或loadTensorBuffer时，报错：The shape of a RGB image should be (h, w, c) or (1, h, w, c), and channels representing R, G, B in order. The provided image shape is [1, 224, 224, 1]
• 调整输出形状为[1,224,224,1]后依然报错
• 直接转Image.memory触发无效图像数据异常
• 用ImageConversions.convertGrayscaleTensorBufferToImage也失败

原代码如下：

ImageProcessor imageProcessor = ImageProcessorBuilder()
    .add(ResizeOp(224, 224, ResizeMethod.NEAREST_NEIGHBOUR))
    .add(NormalizeOp(127.5, 127.5))
    .build();

SequentialProcessor<TensorBuffer> probabilityProcessor = TensorProcessorBuilder().add(DequantizeOp(0, 1 / 255)).build();

TensorImage tensorImage = TensorImage(TfLiteType.float32);
tensorImage.loadImage(img.Image.fromBytes(224, 224, image.readAsBytesSync()));
tensorImage = imageProcessor.process(tensorImage);

TensorBuffer tensorBuffer;

try{
  Interpreter interpreter = await Interpreter.fromAsset('models/enet.tflite');
  tensorBuffer = TensorBuffer.createFixedSize(interpreter.getOutputTensor(0).shape, interpreter.getOutputTensor(0).type);

  interpreter.run(tensorImage.buffer, tensorBuffer.getBuffer());

  tensorBuffer = probabilityProcessor.process(tensorBuffer);

  // ignore: invalid_use_of_protected_member
  tensorBuffer.resize(List<int>.of([1, 224, 224, 1]));

  TensorImage resultImage = TensorImage(tensorBuffer.getDataType());
  resultImage.loadTensorBuffer(tensorBuffer);

}catch(e){
  print('Error loading model: ' + e.toString());
}

尝试过的其他转换代码：

Image result = Image.memory(tensorBuffer.getBuffer().asUint8List());

img.Image resultImage = ImageConversions.convertGrayscaleTensorBufferToImage(tensorBuffer);

问题核心在于分割模型的输出是单通道的类别预测图（每个像素对应类别索引或概率），不是常规的RGB图像，直接用TensorImage或Image.memory转换会因为格式不匹配失败。需要先处理输出张量，将类别值映射为可显示的像素数据，再构建Image。

步骤1：调整张量形状并提取数据 #

模型输出[1,1,224,224]包含冗余的批量和通道维度，先提取出[224,224]的二维数据：

// 提取224x224的分割数据
List<double> outputData = tensorBuffer.getDoubleList();
List<List<double>> segmentationMap = [];
for (int i = 0; i < 224; i++) {
  segmentationMap.add(outputData.sublist(i*224, (i+1)*224));
}

步骤2：将类别值映射为RGB颜色 #

给每个分割类别分配对应的RGB颜色，把类别值转换为可显示的像素数据：

// 自定义类别颜色映射（根据你的模型类别数调整）
Map<int, Color> classColors = {
  0: Color(0xFF000000), // 背景
  1: Color(0xFFFF0000), // 类别1
  2: Color(0xFF00FF00), // 类别2
  3: Color(0xFF0000FF), // 类别3
};

// 构建RGBA格式的像素数组
List<int> pixels = [];
for (int y = 0; y < 224; y++) {
  for (int x = 0; x < 224; x++) {
    // 取类别索引（如果输出是概率分布，需先取argmax）
    int classIndex = segmentationMap[y][x].round();
    Color color = classColors[classIndex] ?? Color(0xFFFFFFFF);
    pixels.addAll([color.red, color.green, color.blue, color.alpha]);
  }
}

步骤3：构建可显示的Image #

用img.Image库构建图像对象，再转换为Flutter可渲染的Image组件：

// 生成img.Image对象
img.Image segmentedImage = img.Image.fromBytes(224, 224, pixels, format: img.Format.rgba);
// 转成Flutter Image
Image displayImage = Image.memory(img.encodePng(segmentedImage));

完整修改后的代码 #

ImageProcessor imageProcessor = ImageProcessorBuilder()
    .add(ResizeOp(224, 224, ResizeMethod.NEAREST_NEIGHBOUR))
    .add(NormalizeOp(127.5, 127.5))
    .build();

SequentialProcessor<TensorBuffer> probabilityProcessor = TensorProcessorBuilder().add(DequantizeOp(0, 1 / 255)).build();

TensorImage tensorImage = TensorImage(TfLiteType.float32);
tensorImage.loadImage(img.Image.fromBytes(224, 224, image.readAsBytesSync()));
tensorImage = imageProcessor.process(tensorImage);

TensorBuffer tensorBuffer;
Image? displayImage;

try{
  Interpreter interpreter = await Interpreter.fromAsset('models/enet.tflite');
  tensorBuffer = TensorBuffer.createFixedSize(interpreter.getOutputTensor(0).shape, interpreter.getOutputTensor(0).type);

  interpreter.run(tensorImage.buffer, tensorBuffer.getBuffer());

  tensorBuffer = probabilityProcessor.process(tensorBuffer);

  // 提取224x224的分割数据
  List<double> outputData = tensorBuffer.getDoubleList();
  List<List<double>> segmentationMap = [];
  for (int i = 0; i < 224; i++) {
    segmentationMap.add(outputData.sublist(i*224, (i+1)*224));
  }

  // 类别颜色映射（根据你的模型类别调整）
  Map<int, Color> classColors = {
    0: Color(0xFF000000),
    1: Color(0xFFFF0000),
    2: Color(0xFF00FF00),
    3: Color(0xFF0000FF),
  };

  // 构建像素数组
  List<int> pixels = [];
  for (int y = 0; y < 224; y++) {
    for (int x = 0; x < 224; x++) {
      // 若输出是概率分布，需先对每个像素的概率数组取argmax得到类别索引
      int classIndex = segmentationMap[y][x].round();
      Color color = classColors[classIndex] ?? Color(0xFFFFFFFF);
      pixels.addAll([color.red, color.green, color.blue, color.alpha]);
    }
  }

  // 生成可显示的Image
  img.Image segmentedImage = img.Image.fromBytes(224, 224, pixels, format: img.Format.rgba);
  displayImage = Image.memory(img.encodePng(segmentedImage));

}catch(e){
  print('Error loading model: ' + e.toString());
}

注意事项 #

• 如果模型输出是概率分布（每个像素对应多个类别的概率），需先对每个像素的概率数组取argmax，得到对应的类别索引后再映射颜色。
• 类别颜色映射要和模型训练时的类别一一对应，确保不同类别颜色区分度足够。
• 避免直接修改TensorBuffer的shape，手动提取数据处理可防止破坏张量内部内存结构。

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