c++ - 具有 1 个通道分段张量的分段应用程序

Question

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

Answer 1

我知道这有点晚了，但如果你还没有弄清楚，可能会对你有所帮助。
由于 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_templateinput_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
    }
  }
}

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