具有 1 个通道分割张量的分割应用程序

如何解决具有 1 个通道分割张量的分割应用程序

我正在尝试使用 mediapipe 的分割示例和我自己的 tflite 模型进行分割。但我的输出是一个单通道分割张量。如何更改 TfLiteTensorsToSegmentationcalculator 以使其适用于我的模型？

解决方法

我知道现在有点晚了，但如果您还没有弄清楚，可能会帮助您。
由于 TfLiteTensorsToSegmentationCalculator 是为套件 2 通道输出而构建的，因此如果要将其更改为 1 通道输出，则需要对 LoadOptions 、processGpu 和相应的计算着色器代码进行更改。注意：我假设您仅使用 GPU 管道而不使用 CPU。首先，您需要在 LoadOptions 函数中允许 1 个通道，如下所示。

mediapipe::Status TfLiteTensorsToSegmentationCalculator::LoadOptions(
    CalculatorContext* cc) {
....
RET_CHECK_EQ(tensor_channels_,2)// Change it to 1

....

接下来您需要更改计算着色器代码。在 shader_src_template input_data 中是一个只读缓冲区，它保存来自 tflite 模型的输出张量。此缓冲区是为 vec2 定义的，您需要将其更改为 vec1。

layout(std430,binding = 2) readonly buffer B0 {
  vec2 elements[];
} input_data;   // data tensor,Change it to vec1

现在在着色器代码的主函数中，您可以看到它们正在使用 input_value.rg，在您的情况下，您只需要使用 input_value.r。您将无法使用 softmax，因为它是一个通道，因此无论您想做什么，都可以对一个通道进行处理，并将结果存储到 R&A 通道中的 out_value。这样您就不需要对下一个计算器进行任何更改，重新着色计算器可以使用与在 2 通道情况下相同的方式使用掩码。
您还需要在 hair_segmentation_mobile_gpu.pbtxt

处更改 TfLiteTensorsToSegmentationCalculatorOptions 文件中的频道数

node {
  calculator: "TfLiteTensorsToSegmentationCalculator"
  input_stream: "TENSORS_GPU:segmentation_tensor"
  input_stream: "PREV_MASK_GPU:previous_hair_mask"
  output_stream: "MASK_GPU:hair_mask"
  node_options: {
    [type.googleapis.com/mediapipe.TfLiteTensorsToSegmentationCalculatorOptions] {
      tensor_width: 512 
      tensor_height: 512
      tensor_channels: 2 # change it to 1
      combine_with_previous_ratio: 0.9
      output_layer_index: 1
    }
  }
}

如果您有任何问题，请告诉我。