async segmentMultiPerson()

in body-pix/src/body_pix_model.ts [597:684]
79 lines of code
2 McCabe index (conditional complexity)

  async segmentMultiPerson(
      input: BodyPixInput,
      config: MultiPersonInstanceInferenceConfig =
          MULTI_PERSON_INSTANCE_INFERENCE_CONFIG):
      Promise<PersonSegmentation[]> {
    config = {...MULTI_PERSON_INSTANCE_INFERENCE_CONFIG, ...config};
    validateMultiPersonInstanceInferenceConfig(config);
    const [height, width] = getInputSize(input);
    const internalResolutionHeightAndWidth = toInputResolutionHeightAndWidth(
        config.internalResolution, this.baseModel.outputStride,
        [height, width]);

    const {resized, padding} =
        padAndResizeTo(input, internalResolutionHeightAndWidth);
    const {
      segmentation,
      longOffsets,
      heatmapScoresRaw,
      offsetsRaw,
      displacementFwdRaw,
      displacementBwdRaw,
    } = tf.tidy(() => {
      const {
        segmentLogits,
        longOffsets,
        heatmapScores,
        offsets,
        displacementFwd,
        displacementBwd,
      } = this.predictForMultiPersonInstanceSegmentationAndPart(resized);
      const scaledSegmentScores = scaleAndCropToInputTensorShape(
          segmentLogits, [height, width], internalResolutionHeightAndWidth,
          [[padding.top, padding.bottom], [padding.left, padding.right]],
          APPLY_SIGMOID_ACTIVATION);
      const longOffsetsResized = false;
      let scaledLongOffsets;
      if (longOffsetsResized) {
        scaledLongOffsets = scaleAndCropToInputTensorShape(
            longOffsets, [height, width], internalResolutionHeightAndWidth,
            [[padding.top, padding.bottom], [padding.left, padding.right]],
            APPLY_SIGMOID_ACTIVATION);
      } else {
        scaledLongOffsets = longOffsets;
      }

      const segmentation = toMaskTensor(
          tf.squeeze(scaledSegmentScores), config.segmentationThreshold);

      return {
        segmentation,
        longOffsets: scaledLongOffsets,
        heatmapScoresRaw: heatmapScores,
        offsetsRaw: offsets,
        displacementFwdRaw: displacementFwd,
        displacementBwdRaw: displacementBwd,
      };
    });

    const tensorBuffers = await toTensorBuffers3D(
        [heatmapScoresRaw, offsetsRaw, displacementFwdRaw, displacementBwdRaw]);
    const [scoresBuf, offsetsBuf, displacementsFwdBuf, displacementsBwdBuf] =
        tensorBuffers;

    let poses = decodeMultiplePoses(
        scoresBuf, offsetsBuf, displacementsFwdBuf, displacementsBwdBuf,
        this.baseModel.outputStride, config.maxDetections,
        config.scoreThreshold, config.nmsRadius);

    poses = scaleAndFlipPoses(
        poses, [height, width], internalResolutionHeightAndWidth, padding,
        FLIP_POSES_AFTER_SCALING);

    const instanceMasks = await decodePersonInstanceMasks(
        segmentation, longOffsets, poses, height, width,
        this.baseModel.outputStride, internalResolutionHeightAndWidth, padding,
        config.scoreThreshold, config.refineSteps, config.minKeypointScore,
        config.maxDetections);

    resized.dispose();
    segmentation.dispose();
    longOffsets.dispose();
    heatmapScoresRaw.dispose();
    offsetsRaw.dispose();
    displacementFwdRaw.dispose();
    displacementBwdRaw.dispose();

    return instanceMasks;
  }