Add keypoint Extraction Post-Processing Layer to TensorFlow Model

I have PoseNet Tensorflow saved model that takes in an image and outputs heatmap and offset tensors. the model works fine, but I just want to add layer to it that performs post-processing.

Currently, I'm extracting the final keypoints in Python code. How can I add post-processing layer to the model itself so it outputs the keypoints along with the scores?

There are models, such as Movenet, that output the final keypoint, and I want to do the same thing for PoseNet.

This image illustrates what I'm trying to accomplish:

I've looked at the following posts about adding post-processing layer, but I don't know how to apply it for my problem:

1. How to add post-processing into a Tensorflow Model?
2. Cannot add layers to saved Keras Model. 'Model' object has no attribute 'add'
3. How to add another layer on a pre-loaded network?

I fully understand the post-processing algorithm and have implemented it in Python. Now, I want to integrate this functionality directly into the model itself:

"""heatmap shape [9, 9, 17]offset shape [9, 9, 34]"""def parse_output(heatmap, offset):    # Get the number of joints - value is 17 for Posenet    joint_num = heatmap.shape[-1]    # Initialize an array to store the keypoints    pose_kps = np.zeros((joint_num, 3), np.uint32)    # Iterate over each joint    for i in range(heatmap.shape[-1]):        # select heatmap for the i-th joint            joint_heatmap = heatmap[..., i]        # Find the maximum probability and its position        max_prob = np.max(joint_heatmap)        # get the x, y coordinates of the max_prob position. eg: [4,7]        max_val_pos = np.squeeze(np.argwhere(joint_heatmap == max_prob))        # this helps to scale the values back to the input image coordinates        remap = np.array(max_val_pos/8*257, dtype=np.int32)        # Assign the calculated values to the keypoints array        pose_kps[i, 0] = remap[0] + offset[max_val_pos[0], max_val_pos[1], i]        pose_kps[i, 1] = remap[1] + offset[max_val_pos[0], max_val_pos[1], i + joint_num]        pose_kps[i, 2] = max_prob    return pose_kps

The above parse_output make sense to me and similar implementation it's also done in the following projects:

1. How to parse the heatmap output for the pose estimation tflite model?
2. decodeSinglePose

I have created a sample PosetNetDemo project to show my current implementation.

I appreciate it if you point me to a resource or help me solve it. Thank you!