Add Landmark 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 have created a sample PosetNetDemo project that shows my current implementation here.