c++ - Segmentation app with 1 channel segmentation tensor

Question

I'm trying to use the segmentation example of mediapipe with my own tflite model for segmentation. But my output is a one channel segmentation tensor. How could I change the TfLiteTensorsToSegmentationCalculator to make it work with my model ?

Answer 1

I know it's bit late but may help you if you haven't figure out yet.
As TfLiteTensorsToSegmentationCalculator is build to suite 2 channel output so if you want to change it to 1 one channel then you need to make changes to LoadOptions , processGpu and corresponding compute shader code. Note: I am assuming you are using only GPU pipeline and not CPU. First change you need to allow 1 channel in LoadOptions function as follows.

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

....

Next you need to change compute shader code. In the shader_src_template input_data is a read only buffer which hold the output tensor from tflite model. This buffer is define for vec2 which you need to change to vec1.

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

Now in the main function of shader code you can see they are using input_value.rg in your case you need to use input_value.r only. You won't be able to use softmax as it is one channel so do your processing for one channel whatever you want to do and store your result to out_value in both R & A channel. This way you don't need to make any changes to next calculator, Recolor calculator can use mask same way it was using in 2 channel case.
Also you need to change number of channel in hair_segmentation_mobile_gpu.pbtxt file at 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
    }
  }
}

Let me know if you have any issue.