import math from typing import ( List, Tuple, ) import cv2 import matplotlib.cm import numpy as np from scipy.ndimage.filters import ( gaussian_filter, maximum_filter, ) from scipy.ndimage.morphology import generate_binary_structure # It is better to use 0.1 as threshold when evaluation, but 0.3 for demo # purpose. cmap = matplotlib.cm.get_cmap("hsv") # Heatmap indices to find each limb (joint connection). Eg: limb_type=1 is # Neck->LShoulder, so joint_to_limb_heatmap_relationship[1] represents the # indices of heatmaps to look for joints: neck=1, LShoulder=5 joint_to_limb_heatmap_relationship = [ [1, 2], [1, 5], [2, 3], [3, 4], [5, 6], [6, 7], [1, 8], [8, 9], [9, 10], [1, 11], [11, 12], [12, 13], [1, 0], [0, 14], [14, 16], [0, 15], [15, 17], [2, 16], [5, 17], ] # PAF indices containing the x and y coordinates of the PAF for a given limb. # Eg: limb_type=1 is Neck->LShoulder, so # PAFneckLShoulder_x=paf_xy_coords_per_limb[1][0] and # PAFneckLShoulder_y=paf_xy_coords_per_limb[1][1] paf_xy_coords_per_limb = [ [12, 13], [20, 21], [14, 15], [16, 17], [22, 23], [24, 25], [0, 1], [2, 3], [4, 5], [6, 7], [8, 9], [10, 11], [28, 29], [30, 31], [34, 35], [32, 33], [36, 37], [18, 19], [26, 27], ] # Color code used to plot different joints and limbs (eg: joint_type=3 and # limb_type=3 will use colors[3]) colors = [ [255, 0, 0], [255, 85, 0], [255, 170, 0], [255, 255, 0], [170, 255, 0], [85, 255, 0], [0, 255, 0], [0, 255, 85], [0, 255, 170], [0, 255, 255], [0, 170, 255], [0, 85, 255], [0, 0, 255], [85, 0, 255], [170, 0, 255], [255, 0, 255], [255, 0, 170], [255, 0, 85], [255, 0, 0], ] NUM_JOINTS = 18 NUM_LIMBS = len(joint_to_limb_heatmap_relationship) def find_peaks(param: dict, img: np.ndarray) -> np.ndarray: """ Finds local maxima in a (grayscale) image whose values are above a given threshold. Args: param (dict): A dictionary containing parameters. - 'thre1' (float): Threshold for peak detection. img (np.ndarray): Input grayscale image (2D array) where peaks are to be found. Returns: np.ndarray: A 2D array containing the [x, y] coordinates of each peak found in the image. """ peaks_binary = (maximum_filter(img, footprint=generate_binary_structure(2, 1)) == img) * ( img > param["thre1"] ) # Note reverse ([::-1]): we return [[x y], [x y]...] instead of [[y x], [y x]...] return np.array(np.nonzero(peaks_binary)[::-1]).T def compute_resized_coords(coords: Tuple[float, float], resizeFactor: float) -> np.ndarray: """ Computes the new coordinates of a cell in an array after resizing the array. Args: coords (Tuple[float, float]): Coordinates (indices) of a cell in some input array. resizeFactor (float): Resize coefficient, indicating how much bigger the destination array is compared to the original one. Returns: np.ndarray: Coordinates in an array of size `shape_dest`, expressing the array indices of the closest point to 'coords' if an image of size `shape_source` was resized to `shape_dest`. """ # 1) Add 0.5 to coords to get coordinates of center of the pixel (e.g. # index [0,0] represents the pixel at location [0.5,0.5]) # 2) Transform those coordinates to shape_dest, by multiplying by resizeFactor # 3) That number represents the location of the pixel center in the new array, # so subtract 0.5 to get coordinates of the array index/indices (revert step 1) return (np.array(coords, dtype=float) + 0.5) * resizeFactor - 0.5 def NMS( param: dict, heatmaps: np.ndarray, upsampFactor: float = 1.0, bool_refine_center: bool = True, bool_gaussian_filt: bool = False, ) -> List[np.ndarray]: """ Performs Non-Maxima Suppression (NMS) to find peaks (local maxima) in a set of grayscale images. Args: param (dict): Additional parameters for NMS. heatmaps (np.ndarray): Set of grayscale images on which to find local maxima. A 3D numpy array with dimensions image_height x image_width x num_heatmaps. upsampFactor (float): Size ratio between CPM (Convolutional Pose Machine) heatmap output and the input image size. For example, upsampFactor=16 if the original image was 480x640 and heatmaps are 30x40xN. bool_refine_center (bool): Flag indicating whether to refine the center of the peak. Defaults to True. If True, the function upsamples a small patch around each low-res peak and fine-tunes the location of the peak at the resolution of the original input image. If False, simply returns the low-res peak found upscaled by upsampFactor. bool_gaussian_filt (bool): Flag indicating whether to apply a 1D Gaussian filter (smoothing) to each upsampled patch before fine-tuning the location of each peak. Defaults to False. Returns: np.ndarray: A NUM_JOINTS x 4 numpy array where each row represents a joint type (0=nose, 1=neck...) and the columns indicate the {x,y} position, the score (probability), and a unique id (counter). """ joint_list_per_joint_type = [] cnt_total_joints = 0 # For every peak found, win_size specifies how many pixels in each # direction from the peak we take to obtain the patch that will be # upsampled. Eg: win_size=1 -> patch is 3x3; win_size=2 -> 5x5 # (for BICUBIC interpolation to be accurate, win_size needs to be >=2!) win_size = 2 for joint in range(NUM_JOINTS): map_orig = heatmaps[:, :, joint] peak_coords = find_peaks(param, map_orig) peaks = np.zeros((len(peak_coords), 4)) for i, peak in enumerate(peak_coords): if bool_refine_center: x_min, y_min = np.maximum(0, peak - win_size) x_max, y_max = np.minimum(np.array(map_orig.T.shape) - 1, peak + win_size) # Take a small patch around each peak and only upsample that tiny region patch = map_orig[y_min : y_max + 1, x_min : x_max + 1] map_upsamp = cv2.resize( patch, None, fx=upsampFactor, fy=upsampFactor, interpolation=cv2.INTER_CUBIC ) # Gaussian filtering takes an average of 0.8ms/peak (and there might be # more than one peak per joint!) -> For now, skip it (it's # accurate enough) map_upsamp = ( gaussian_filter(map_upsamp, sigma=3) if bool_gaussian_filt else map_upsamp ) # Obtain the coordinates of the maximum value in the patch location_of_max = np.unravel_index(map_upsamp.argmax(), map_upsamp.shape) # Remember that peaks indicates [x,y] -> need to reverse it for # [y,x] location_of_patch_center = compute_resized_coords( peak[::-1] - [y_min, x_min], upsampFactor ) # Calculate the offset wrt to the patch center where the actual # maximum is refined_center = location_of_max - location_of_patch_center peak_score = map_upsamp[location_of_max] else: refined_center = [0, 0] # Flip peak coordinates since they are [x,y] instead of [y,x] peak_score = map_orig[tuple(peak[::-1])] peaks[i, :] = tuple( [ int(round(x)) for x in compute_resized_coords(peak_coords[i], upsampFactor) + refined_center[::-1] ] ) + (peak_score, cnt_total_joints) cnt_total_joints += 1 joint_list_per_joint_type.append(peaks) return joint_list_per_joint_type def find_connected_joints( param: dict, paf_upsamp: np.ndarray, joint_list_per_joint_type: List[np.ndarray], num_intermed_pts: int = 10, ) -> List[np.ndarray]: """ For every type of limb (e.g., forearm, shin, etc.), looks for every potential pair of joints (e.g., every wrist-elbow combination) and evaluates the PAFs to determine which pairs are indeed body limbs. Args: paf_upsamp (np.ndarray): PAFs upsampled to the original input image resolution. joint_list_per_joint_type (List[np.ndarray]): List of joint lists per joint type. See the 'return' doc of NMS(). num_intermed_pts (int): Number of intermediate points to take between joint_src and joint_dst, at which the PAFs will be evaluated. Defaults to 10. Returns: List[np.ndarray]: List of NUM_LIMBS rows. For every limb_type (a row) we store a list of all limbs of that type found (e.g., all the right forearms). For each limb (each item in connected_limbs[limb_type]), we store 5 cells: {joint_src_id, joint_dst_id}: a unique number associated with each joint, limb_score_penalizing_long_dist: a score of how good a connection of the joints is, penalized if the limb length is too long, {joint_src_index, joint_dst_index}: the index of the joint within all the joints of that type found (e.g., the 3rd right elbow found) """ connected_limbs = [] # Auxiliary array to access paf_upsamp quickly limb_intermed_coords = np.empty((4, num_intermed_pts), dtype=np.intp) for limb_type in range(NUM_LIMBS): # List of all joints of type A found, where A is specified by limb_type # (eg: a right forearm starts in a right elbow) joints_src = joint_list_per_joint_type[joint_to_limb_heatmap_relationship[limb_type][0]] # List of all joints of type B found, where B is specified by limb_type # (eg: a right forearm ends in a right wrist) joints_dst = joint_list_per_joint_type[joint_to_limb_heatmap_relationship[limb_type][1]] if len(joints_src) == 0 or len(joints_dst) == 0: # No limbs of this type found (eg: no right forearms found because # we didn't find any right wrists or right elbows) connected_limbs.append([]) else: connection_candidates = [] # Specify the paf index that contains the x-coord of the paf for # this limb limb_intermed_coords[2, :] = paf_xy_coords_per_limb[limb_type][0] # And the y-coord paf index limb_intermed_coords[3, :] = paf_xy_coords_per_limb[limb_type][1] for i, joint_src in enumerate(joints_src): # Try every possible joints_src[i]-joints_dst[j] pair and see # if it's a feasible limb for j, joint_dst in enumerate(joints_dst): # Subtract the position of both joints to obtain the # direction of the potential limb limb_dir = joint_dst[:2] - joint_src[:2] # Compute the distance/length of the potential limb (norm # of limb_dir) limb_dist = np.sqrt(np.sum(limb_dir**2)) + 1e-8 limb_dir = limb_dir / limb_dist # Normalize limb_dir to be a unit vector # Linearly distribute num_intermed_pts points from the x # coordinate of joint_src to the x coordinate of joint_dst limb_intermed_coords[1, :] = np.round( np.linspace(joint_src[0], joint_dst[0], num=num_intermed_pts) ) limb_intermed_coords[0, :] = np.round( np.linspace(joint_src[1], joint_dst[1], num=num_intermed_pts) ) # Same for the y coordinate intermed_paf = paf_upsamp[ limb_intermed_coords[0, :], limb_intermed_coords[1, :], limb_intermed_coords[2:4, :], ].T score_intermed_pts = intermed_paf.dot(limb_dir) score_penalizing_long_dist = score_intermed_pts.mean() + min( 0.5 * paf_upsamp.shape[0] / limb_dist - 1, 0 ) # Criterion 1: At least 80% of the intermediate points have # a score higher than thre2 criterion1 = ( np.count_nonzero(score_intermed_pts > param["thre2"]) > 0.8 * num_intermed_pts ) # Criterion 2: Mean score, penalized for large limb # distances (larger than half the image height), is # positive criterion2 = score_penalizing_long_dist > 0 if criterion1 and criterion2: # Last value is the combined paf(+limb_dist) + heatmap # scores of both joints connection_candidates.append( [ i, j, score_penalizing_long_dist, score_penalizing_long_dist + joint_src[2] + joint_dst[2], ] ) # Sort connection candidates based on their # score_penalizing_long_dist connection_candidates = sorted(connection_candidates, key=lambda x: x[2], reverse=True) connections = np.empty((0, 5)) # There can only be as many limbs as the smallest number of source # or destination joints (eg: only 2 forearms if there's 5 wrists # but 2 elbows) max_connections = min(len(joints_src), len(joints_dst)) # Traverse all potential joint connections (sorted by their score) for potential_connection in connection_candidates: i, j, s = potential_connection[0:3] # Make sure joints_src[i] or joints_dst[j] haven't already been # connected to other joints_dst or joints_src if i not in connections[:, 3] and j not in connections[:, 4]: # [joint_src_id, joint_dst_id, limb_score_penalizing_long_dist, joint_src_index, joint_dst_index] connections = np.vstack( [connections, [joints_src[i][3], joints_dst[j][3], s, i, j]] ) # Exit if we've already established max_connections # connections (each joint can't be connected to more than # one joint) if len(connections) >= max_connections: break connected_limbs.append(connections) return connected_limbs def group_limbs_of_same_person( connected_limbs: List[np.ndarray], joint_list: np.ndarray ) -> np.ndarray: """ Associate limbs belonging to the same person together. Args: connected_limbs (List[np.ndarray]): List of connected limbs. See the 'return' doc of find_connected_joints(). joint_list (np.ndarray): Unraveled version of joint_list_per_joint. See the 'return' doc of NMS(). Returns: np.ndarray: A 2D array of size num_people x (NUM_JOINTS+2). For each person found: - First NUM_JOINTS columns contain the index (in joint_list) of the joints associated with that person (or -1 if their i-th joint wasn't found). - 2nd-to-last column: Overall score of the joints+limbs that belong to this person. """ person_to_joint_assoc = [] for limb_type in range(NUM_LIMBS): joint_src_type, joint_dst_type = joint_to_limb_heatmap_relationship[limb_type] for limb_info in connected_limbs[limb_type]: person_assoc_idx = [] for person, person_limbs in enumerate(person_to_joint_assoc): if ( person_limbs[joint_src_type] == limb_info[0] or person_limbs[joint_dst_type] == limb_info[1] ): person_assoc_idx.append(person) # If one of the joints has been associated to a person, and either # the other joint is also associated with the same person or not # associated to anyone yet: if len(person_assoc_idx) == 1: person_limbs = person_to_joint_assoc[person_assoc_idx[0]] # If the other joint is not associated to anyone yet, if person_limbs[joint_dst_type] != limb_info[1]: # Associate it with the current person person_limbs[joint_dst_type] = limb_info[1] # Increase the number of limbs associated to this person person_limbs[-1] += 1 # And update the total score (+= heatmap score of joint_dst # + score of connecting joint_src with joint_dst) person_limbs[-2] += joint_list[limb_info[1].astype(int), 2] + limb_info[2] elif len(person_assoc_idx) == 2: # if found 2 and disjoint, merge them person1_limbs = person_to_joint_assoc[person_assoc_idx[0]] person2_limbs = person_to_joint_assoc[person_assoc_idx[1]] membership = ((person1_limbs >= 0) & (person2_limbs >= 0))[:-2] if ( not membership.any() ): # If both people have no same joints connected, merge them into a single person # Update which joints are connected person1_limbs[:-2] += person2_limbs[:-2] + 1 # Update the overall score and total count of joints # connected by summing their counters person1_limbs[-2:] += person2_limbs[-2:] # Add the score of the current joint connection to the # overall score person1_limbs[-2] += limb_info[2] person_to_joint_assoc.pop(person_assoc_idx[1]) else: # Same case as len(person_assoc_idx)==1 above person1_limbs[joint_dst_type] = limb_info[1] person1_limbs[-1] += 1 person1_limbs[-2] += joint_list[limb_info[1].astype(int), 2] + limb_info[2] else: # No person has claimed any of these joints, create a new person # Initialize person info to all -1 (no joint associations) row = -1 * np.ones(20) # Store the joint info of the new connection row[joint_src_type] = limb_info[0] row[joint_dst_type] = limb_info[1] # Total count of connected joints for this person: 2 row[-1] = 2 # Compute overall score: score joint_src + score joint_dst + score connection # {joint_src,joint_dst} row[-2] = sum(joint_list[limb_info[:2].astype(int), 2]) + limb_info[2] person_to_joint_assoc.append(row) # Delete people who have very few parts connected people_to_delete = [] for person_id, person_info in enumerate(person_to_joint_assoc): if person_info[-1] < 3 or person_info[-2] / person_info[-1] < 0.2: people_to_delete.append(person_id) # Traverse the list in reverse order so we delete indices starting from the # last one (otherwise, removing item for example 0 would modify the indices of # the remaining people to be deleted!) for index in people_to_delete[::-1]: person_to_joint_assoc.pop(index) # Appending items to a np.array can be very costly (allocating new memory, copying over the array, then adding new row) # Instead, we treat the set of people as a list (fast to append items) and # only convert to np.array at the end return np.array(person_to_joint_assoc) def plot_pose( img_orig: np.ndarray, joint_list: np.ndarray, person_to_joint_assoc: np.ndarray, bool_fast_plot: bool = True, plot_ear_to_shoulder: bool = False, ) -> Tuple[np.ndarray, np.ndarray]: canvas = img_orig.copy() # Make a copy so we don't modify the original image # to_plot is the location of all joints found overlaid on top of the # original image to_plot = canvas.copy() if bool_fast_plot else cv2.addWeighted(img_orig, 0.3, canvas, 0.7, 0) limb_thickness = 4 # Last 2 limbs connect ears with shoulders and this looks very weird. # Disabled by default to be consistent with original rtpose output which_limbs_to_plot = NUM_LIMBS if plot_ear_to_shoulder else NUM_LIMBS - 2 for limb_type in range(which_limbs_to_plot): for person_joint_info in person_to_joint_assoc: joint_indices = person_joint_info[ joint_to_limb_heatmap_relationship[limb_type] ].astype(int) if -1 in joint_indices: # Only draw actual limbs (connected joints), skip if not # connected continue # joint_coords[:,0] represents Y coords of both joints; # joint_coords[:,1], X coords joint_coords = joint_list[joint_indices, 0:2] for joint in joint_coords: cv2.circle(canvas, tuple(joint[0:2].astype(int)), 2, (255, 255, 255), thickness=-1) # mean along the axis=0 computes meanYcoord and meanXcoord -> Round # and make int to avoid errors coords_center = tuple(np.round(np.mean(joint_coords, 0)).astype(int)) # joint_coords[0,:] is the coords of joint_src; joint_coords[1,:] # is the coords of joint_dst limb_dir = joint_coords[0, :] - joint_coords[1, :] limb_length = np.linalg.norm(limb_dir) # Get the angle of limb_dir in degrees using atan2(limb_dir_x, limb_dir_y) angle = math.degrees(math.atan2(limb_dir[1], limb_dir[0])) # For faster plotting, just plot over canvas instead of constantly # copying it cur_canvas = canvas if bool_fast_plot else canvas.copy() polygon = cv2.ellipse2Poly( coords_center, (int(limb_length / 2), limb_thickness), int(angle), 0, 360, 1 ) cv2.fillConvexPoly(cur_canvas, polygon, colors[limb_type]) if not bool_fast_plot: canvas = cv2.addWeighted(canvas, 0.4, cur_canvas, 0.6, 0) return to_plot, canvas def decode_pose( img_orig: np.ndarray, heatmaps: np.ndarray, pafs: np.ndarray ) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: param = {"thre1": 0.1, "thre2": 0.05, "thre3": 0.5} # Bottom-up approach: # Step 1: find all joints in the image (organized by joint type: [0]=nose, [1]=neck...) joint_list_per_joint_type = NMS(param, heatmaps, img_orig.shape[0] / float(heatmaps.shape[0])) # joint_list is an unravel'd version of joint_list_per_joint, where we add # a 5th column to indicate the joint_type (0=nose, 1=neck...) joint_list = np.array( [ tuple(peak) + (joint_type,) for joint_type, joint_peaks in enumerate(joint_list_per_joint_type) for peak in joint_peaks ] ) # Step 2: find which joints go together to form limbs (which wrists go with which elbows) paf_upsamp = cv2.resize( pafs, (img_orig.shape[1], img_orig.shape[0]), interpolation=cv2.INTER_CUBIC, ) connected_limbs = find_connected_joints(param, paf_upsamp, joint_list_per_joint_type) # Step 3: associate limbs that belong to the same person person_to_joint_assoc = group_limbs_of_same_person(connected_limbs, joint_list) # Step 4: plot results to_plot, canvas = plot_pose(img_orig, joint_list, person_to_joint_assoc) return to_plot, canvas, joint_list, person_to_joint_assoc