Add keypoint Extraction Post-Processing Layer to TensorFlow Model

I have a TensorFlow PoseNet model that takes in an image and outputs heatmap and offset tensors. Currently, I'm extracting the final keypoints in Python code. How can I add a 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.

I've looked at the following posts about adding post-processing layer:

I found a post How to add post-processing into a Tensorflow Model?, but I don't know how to apply it for my problem.

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