Project Code

def extract_keypoints(heatmap, all_keypoints, total_keypoint_num):
    heatmap[heatmap < 0.1] = 0
    heatmap_with_borders = np.pad(heatmap, [(2, 2), (2, 2)], mode='constant')
    heatmap_center = heatmap_with_borders[1:heatmap_with_borders.shape[0]-1, 1:heatmap_with_borders.shape[1]-1]
    heatmap_left = heatmap_with_borders[1:heatmap_with_borders.shape[0]-1, 2:heatmap_with_borders.shape[1]]
    heatmap_right = heatmap_with_borders[1:heatmap_with_borders.shape[0]-1, 0:heatmap_with_borders.shape[1]-2]
    heatmap_up = heatmap_with_borders[2:heatmap_with_borders.shape[0], 1:heatmap_with_borders.shape[1]-1]
    heatmap_down = heatmap_with_borders[0:heatmap_with_borders.shape[0]-2, 1:heatmap_with_borders.shape[1]-1]

    heatmap_peaks = (heatmap_center > heatmap_left) &\\
                    (heatmap_center > heatmap_right) &\\
                    (heatmap_center > heatmap_up) &\\
                    (heatmap_center > heatmap_down)
    heatmap_peaks = heatmap_peaks[1:heatmap_center.shape[0]-1, 1:heatmap_center.shape[1]-1]
    keypoints = list(zip(np.nonzero(heatmap_peaks)[1], np.nonzero(heatmap_peaks)[0]))  # (w, h)
    keypoints = sorted(keypoints, key=itemgetter(0))

    suppressed = np.zeros(len(keypoints), np.uint8)
    keypoints_with_score_and_id = []
    keypoint_num = 0
    for i in range(len(keypoints)):
        if suppressed[i]:
            continue
        for j in range(i+1, len(keypoints)):
            if math.sqrt((keypoints[i][0] - keypoints[j][0]) ** 2 +
                         (keypoints[i][1] - keypoints[j][1]) ** 2) < 6:
                suppressed[j] = 1
        keypoint_with_score_and_id = (keypoints[i][0], keypoints[i][1], heatmap[keypoints[i][1], keypoints[i][0]],
                                      total_keypoint_num + keypoint_num)
        keypoints_with_score_and_id.append(keypoint_with_score_and_id)
        keypoint_num += 1
    all_keypoints.append(keypoints_with_score_and_id)
    return keypoint_num

The extract_keypoints function is responsible for detecting keypoints (e.g., body joints) from a given heatmap. It identifies local peaks in the heatmap, assigns them scores, and associates them with unique IDs. These keypoints are then used in subsequent steps of the pose estimation pipeline.

Function Definition

def extract_keypoints(heatmap, all_keypoints, total_keypoint_num):

• Parameters:
  • heatmap: A 2D array representing the heatmap for a specific keypoint type (e.g., wrist, elbow). Each value indicates the confidence of a keypoint being present at that location.
  • all_keypoints: A list to store all detected keypoints, including their coordinates, scores, and IDs.
  • total_keypoint_num: The total number of keypoints detected so far (used to assign unique IDs).
• Returns:
  • keypoint_num: The number of keypoints detected in the current heatmap.

Step-by-Step Explanation

1. Thresholding the Heatmap:

   heatmap[heatmap < 0.1] = 0
   

   • Any value in the heatmap below 0.1 is set to 0. This removes low-confidence detections and focuses on stronger peaks.

2. Padding the Heatmap:

   heatmap_with_borders = np.pad(heatmap, [(2, 2), (2, 2)], mode='constant')
   

   • The heatmap is padded with a border of zeros. This ensures that peaks near the edges of the heatmap can still be processed without causing indexing errors.

3. Extracting Neighboring Regions:

   heatmap_center = heatmap_with_borders[1:heatmap_with_borders.shape[0]-1, 1:heatmap_with_borders.shape[1]-1]
   heatmap_left = heatmap_with_borders[1:heatmap_with_borders.shape[0]-1, 2:heatmap_with_borders.shape[1]]
   heatmap_right = heatmap_with_borders[1:heatmap_with_borders.shape[0]-1, 0:heatmap_with_borders.shape[1]-2]
   heatmap_up = heatmap_with_borders[2:heatmap_with_borders.shape[0], 1:heatmap_with_borders.shape[1]-1]
   heatmap_down = heatmap_with_borders[0:heatmap_with_borders.shape[0]-2, 1:heatmap_with_borders.shape[1]-1]
   

   • The function extracts the center region of the heatmap and its neighboring regions (left, right, up, down). These are used to identify local peaks.

4. Finding Local Peaks:

   heatmap_peaks = (heatmap_center > heatmap_left) &\\\\
                   (heatmap_center > heatmap_right) &\\\\
                   (heatmap_center > heatmap_up) &\\\\
                   (heatmap_center > heatmap_down)
   

   • A local peak is a point in the heatmap that is greater than its neighbors in all directions (left, right, up, down).
   • The result is a boolean mask (heatmap_peaks) indicating the locations of these peaks.

5. Extracting Keypoint Coordinates:

   keypoints = list(zip(np.nonzero(heatmap_peaks)[1], np.nonzero(heatmap_peaks)[0]))  # (w, h)
   keypoints = sorted(keypoints, key=itemgetter(0))
   

   • The coordinates of the peaks are extracted using np.nonzero.
   • The keypoints are sorted by their x-coordinate (width) for consistency.

6. Non-Maximum Suppression:

   suppressed = np.zeros(len(keypoints), np.uint8)
   keypoints_with_score_and_id = []
   keypoint_num = 0
   for i in range(len(keypoints)):
       if suppressed[i]:
           continue
       for j in range(i+1, len(keypoints)):
           if math.sqrt((keypoints[i][0] - keypoints[j][0]) ** 2 +
                        (keypoints[i][1] - keypoints[j][1]) ** 2) < 6:
               suppressed[j] = 1
   

   • Non-maximum suppression is applied to remove duplicate or closely located keypoints.
   • If two keypoints are within a distance of 6 pixels, only the first one is kept.

7. Assigning Scores and IDs:

   keypoint_with_score_and_id = (keypoints[i][0], keypoints[i][1], heatmap[keypoints[i][1], keypoints[i][0]],
                                 total_keypoint_num + keypoint_num)
   keypoints_with_score_and_id.append(keypoint_with_score_and_id)
   keypoint_num += 1
   

   • For each remaining keypoint:
     • Its coordinates (x, y) are stored.
     • Its confidence score is taken from the heatmap.
     • A unique ID is assigned based on total_keypoint_num.

8. Storing Keypoints:

   all_keypoints.append(keypoints_with_score_and_id)
   

   • The detected keypoints (with coordinates, scores, and IDs) are added to the all_keypoints list.

9. Returning the Number of Keypoints:

   return keypoint_num
   

   • The function returns the number of keypoints detected in the current heatmap.