Flutter中TensorFlow Lite图像分割仅输出4类中的2类问题排查

问题概述 #

使用TensorFlow Lite在Flutter中开发图像分割项目，模型目标是将图像分割为4个类别，但实际输出仅类别2和3有像素，类别0、1的概率值极低且未被归类。从日志可见，模型原始输出的像素概率包含4类数值，但经过argmax处理后只有2、3类被选中。

可能的问题点及排查方向 #

1. 输入图像预处理与模型训练流程不匹配 #

代码中仅对输入图像做了ResizeOp(512, 512, ResizeMethod.NEAREST_NEIGHBOUR)尺寸调整，但需确认：

• 模型训练时是否需要归一化/标准化：比如是否将像素值从0-255缩放到0-1或-1到1区间，当前代码未做该操作，可能导致模型输出概率分布异常。
• 模型训练时的resize方法：如果训练阶段使用的是BILINEAR而非NEAREST_NEIGHBOUR，也会影响最终输出结果。

2. Interpreter初始化的输出张量形状/类型错误 #

你定义的输出张量结构为：

final output = {
  0: [List<List<List<double>>>.filled(512, List<List<double>>.filled(512, List<double>.filled(4, 0)))],
};

需验证：

• 模型实际输出张量形状：如果模型输出是[1, 4, 512, 512]（通道在前）而非你假设的[1, 512, 512, 4]，那么rawOutput[y][x]读取到的概率值会对应错误的类别，导致0、1类概率被错误读取。
• 可通过print(interpreter.getOutputTensor(0).shape)打印输出形状，确认是否与定义的结构匹配。

3. Argmax函数的潜在逻辑问题 #

当前argmax实现：

int argmax(List<double> list) {
  return list.indexOf(list.reduce((a, b) => a > b ? a : b));
}

当列表存在多个相同最大值时，indexOf会返回第一个出现的索引。虽然从日志看0、1类概率远低于2、3类，暂时不影响，但可以优化为遍历找索引避免潜在问题：

int argmax(List<double> list) {
  int maxIndex = 0;
  double maxValue = list[0];
  for (int i = 1; i < list.length; i++) {
    if (list[i] > maxValue) {
      maxValue = list[i];
      maxIndex = i;
    }
  }
  return maxIndex;
}

4. 模型本身的训练问题 #

从日志中0、1类的极低概率值（比如0.000976、0.000002）来看，模型本身对这两类的识别置信度极低，可能原因：

• 训练数据集类别分布不均衡：0、1类样本数量过少，导致模型学习不足。
• 训练时损失函数或参数设置不合理，模型偏向于预测样本量多的类别。
• 模型导出为TFLite时的量化或转换错误，导致部分类别信息丢失。

排查步骤建议 #

1. 打印interpreter.getOutputTensor(0).shape确认输出张量形状，确保rawOutput[y][x]对应像素的4类概率值。
2. 对齐模型训练时的预处理流程，在Flutter代码中添加归一化操作，示例：

   final ImageProcessor imageProcessor = ImageProcessorBuilder()
       .add(ResizeOp(512, 512, ResizeMethod.BILINEAR)) // 匹配训练时的resize方法
       .add(NormalizeOp(0, 255)) // 若训练时将像素缩放到0-1区间
       // 若训练时用均值方差归一化则替换为：.add(NormalizeOp(127.5, 127.5))
       .build();
   

3. 在Python的TensorFlow原生环境中测试模型输出，对比Flutter中的结果是否一致，排除TFLite转换的问题。
4. 检查训练数据集的类别分布，确认0、1类有足够样本且标注正确。

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