diff --git a/README.md b/README.md index ee14f8e..f2a8757 100644 --- a/README.md +++ b/README.md @@ -33,6 +33,8 @@ This project contains the following scripts and jupyter notebooks: **test_tflite_model.ipynb** - helper notebook to verify exported *TFLite* model. +**estimation_example/** - This is an example demonstrating the estimation algorithm. Here you will find sample heatmaps and pafs dumped into numpy arrays (*.npy) and some scripts: *coordinates.py*, *connections.py*, *estimators.py* containing the code for each step of the estimation algorithm. You can run these scripts separately to better understand each step. In addition, there is the script: *example.py* that shows all the steps together. This script creates an output image with the connections. + # Installation ## Prerequisites diff --git a/estimation_example/config.py b/estimation_example/config.py new file mode 100644 index 0000000..dddfa0f --- /dev/null +++ b/estimation_example/config.py @@ -0,0 +1,152 @@ +from enum import IntEnum + + +class BodyPart(IntEnum): + """ + List of all body parts + """ + nose = 0 + neck = 1 # this part is not from COCO + right_shoulder = 2 + right_elbow = 3 + right_wrist = 4 + left_shoulder = 5 + left_elbow = 6 + left_wrist = 7 + right_hip = 8 + right_knee = 9 + right_ankle = 10 + left_hip = 11 + left_knee = 12 + left_ankle = 13 + right_eye = 14 + left_eye = 15 + right_ear = 16 + left_ear = 17 + background = 18 + + +class ConnectionMeta: + """ + Metadata for each connection type: + + -first body part identifier, connections are defined beteen 2 body parts + -second body part identifier + -index in paf (dx) + -index in paf (dy) + -color, helpful for rendering this connection + """ + def __init__(self, from_body_part: BodyPart, to_body_part: BodyPart, paf_dx_idx: int, + paf_dy_idx: int, color: list): + self.from_body_part = from_body_part + self.to_body_part = to_body_part + self.paf_dx_idx = paf_dx_idx + self.paf_dy_idx = paf_dy_idx + self.color = color + + +class BodyPartMeta: + """ + Metadata for each body part type: + + -body part identifier + -index in heatmap where the relevant peaks can be found for this body part type + -slot index for this body part. During the estimation phase, each skeleton has an array containing + identifiers of body parts which belong to this skeleton. Each such identifier has to be stored at an + specific position in the array. This position is being kept here as a slot_idx + -color, helpful for rendering this body part + """ + def __init__(self, body_part: BodyPart, heatmap_idx: int, slot_idx: int, color: list): + self.body_part = body_part + self.heatmap_idx = heatmap_idx + self.slot_idx = slot_idx + self.color = color + + +class ConnectionsConfig: + """ + Configuration of all body part types and connection types beetween them. This architecture allows you + to register only a subset of body parts and connections. Less connections, faster estimation. + """ + body_parts = dict() + connection_types = [] + + def __init__(self): + self.slot_idx_seq = 0 + + def register_body_part(self, body_part: BodyPart, heatmap_idx: int, color: list): + """ + Registers a body part + """ + self.body_parts[body_part] = BodyPartMeta(body_part, heatmap_idx, self.slot_idx_seq, color) + self.slot_idx_seq += 1 + + def add_connection(self, from_body_part: BodyPart, to_body_part: BodyPart, paf_dx_idx: int, paf_dy_idx: int, color: list): + """ + Adds a connection definition between two body parts. An Exception will be raise if the body part is not registered + """ + if from_body_part not in self.body_parts.keys(): + raise Exception(f"Body part '{from_body_part.name}' is not registered.") + if to_body_part not in self.body_parts.keys(): + raise Exception(f"Body part '{to_body_part.name}' is not registered.") + self.connection_types.append(ConnectionMeta(from_body_part, to_body_part, paf_dx_idx, paf_dy_idx, color)) + + def conn_types_size(self): + """ + Returns the number of all connection types + """ + return len(self.connection_types) + + def body_parts_size(self): + """ + Returns the number of all registered body parts + """ + return len(self.body_parts) + + +def get_default_configuration(): + """ + This is the default configuration including all body parts and connections. + You may remove the last 2 connections - ears to shoulders. Why did the CMU include them in their solution ? + """ + config = ConnectionsConfig() + config.register_body_part(body_part = BodyPart.nose, heatmap_idx = 0, color = [255, 0, 0]) + config.register_body_part(body_part = BodyPart.neck, heatmap_idx = 1, color = [255, 85, 0]) + config.register_body_part(body_part = BodyPart.right_shoulder, heatmap_idx = 2, color = [255, 170, 0]) + config.register_body_part(body_part = BodyPart.right_elbow, heatmap_idx = 3, color = [255, 255, 0]) + config.register_body_part(body_part = BodyPart.right_wrist, heatmap_idx = 4, color = [170, 255, 0]) + config.register_body_part(body_part = BodyPart.left_shoulder, heatmap_idx = 5, color = [85, 255, 0]) + config.register_body_part(body_part = BodyPart.left_elbow, heatmap_idx = 6, color = [0, 255, 0]) + config.register_body_part(body_part = BodyPart.left_wrist, heatmap_idx = 7, color = [0, 255, 85]) + config.register_body_part(body_part = BodyPart.right_hip, heatmap_idx = 8, color = [0, 255, 170]) + config.register_body_part(body_part = BodyPart.right_knee, heatmap_idx = 9, color = [0, 255, 255]) + config.register_body_part(body_part = BodyPart.right_ankle, heatmap_idx = 10, color = [0, 170, 255]) + config.register_body_part(body_part = BodyPart.left_hip, heatmap_idx = 11, color = [0, 85, 255]) + config.register_body_part(body_part = BodyPart.left_knee, heatmap_idx = 12, color = [0, 0, 255]) + config.register_body_part(body_part = BodyPart.left_ankle, heatmap_idx = 13, color = [170, 0, 255]) + config.register_body_part(body_part = BodyPart.right_eye, heatmap_idx = 14, color = [255, 0, 255]) + config.register_body_part(body_part = BodyPart.left_eye, heatmap_idx = 15, color = [255, 0, 170]) + config.register_body_part(body_part = BodyPart.right_ear, heatmap_idx = 16, color = [255, 0, 85]) + config.register_body_part(body_part = BodyPart.left_ear, heatmap_idx = 17, color = [255, 0, 85]) + + config.add_connection(from_body_part = BodyPart.neck, to_body_part = BodyPart.right_shoulder, paf_dx_idx = 12, paf_dy_idx = 13, color = [255, 0, 0]) + config.add_connection(from_body_part = BodyPart.neck, to_body_part = BodyPart.left_shoulder, paf_dx_idx = 20, paf_dy_idx = 21, color = [255, 85, 0]) + config.add_connection(from_body_part = BodyPart.right_shoulder, to_body_part = BodyPart.right_elbow, paf_dx_idx = 14, paf_dy_idx = 15, color = [255, 170, 0]) + config.add_connection(from_body_part = BodyPart.right_elbow, to_body_part = BodyPart.right_wrist, paf_dx_idx = 16, paf_dy_idx = 17, color = [255, 255, 0]) + config.add_connection(from_body_part = BodyPart.left_shoulder, to_body_part = BodyPart.left_elbow, paf_dx_idx = 22, paf_dy_idx = 23, color = [170, 255, 0]) + config.add_connection(from_body_part = BodyPart.left_elbow, to_body_part = BodyPart.left_wrist, paf_dx_idx = 24, paf_dy_idx = 25, color = [85, 255, 0]) + config.add_connection(from_body_part = BodyPart.neck, to_body_part = BodyPart.right_hip, paf_dx_idx = 0, paf_dy_idx = 1, color = [0, 255, 0]) + config.add_connection(from_body_part = BodyPart.right_hip, to_body_part = BodyPart.right_knee, paf_dx_idx = 2, paf_dy_idx = 3, color = [0, 255, 85]) + config.add_connection(from_body_part = BodyPart.right_knee, to_body_part = BodyPart.right_ankle, paf_dx_idx = 4, paf_dy_idx = 5, color = [0, 255, 170]) + config.add_connection(from_body_part = BodyPart.neck, to_body_part = BodyPart.left_hip, paf_dx_idx = 6, paf_dy_idx = 7, color = [0, 255, 255]) + config.add_connection(from_body_part = BodyPart.left_hip, to_body_part = BodyPart.left_knee, paf_dx_idx = 8, paf_dy_idx = 9, color = [0, 170, 255]) + config.add_connection(from_body_part = BodyPart.left_knee, to_body_part = BodyPart.left_ankle, paf_dx_idx = 10, paf_dy_idx = 11, color = [0, 85, 255]) + config.add_connection(from_body_part = BodyPart.neck, to_body_part = BodyPart.nose, paf_dx_idx = 28, paf_dy_idx = 29, color = [0, 0, 255]) + config.add_connection(from_body_part = BodyPart.nose, to_body_part = BodyPart.right_eye, paf_dx_idx = 30, paf_dy_idx = 31, color = [85, 0, 255]) + config.add_connection(from_body_part = BodyPart.right_eye, to_body_part = BodyPart.right_ear, paf_dx_idx = 34, paf_dy_idx = 35, color = [170, 0, 255]) + config.add_connection(from_body_part = BodyPart.nose, to_body_part = BodyPart.left_eye, paf_dx_idx = 32, paf_dy_idx = 33, color = [255, 0, 255]) + config.add_connection(from_body_part = BodyPart.left_eye, to_body_part = BodyPart.left_ear, paf_dx_idx = 36, paf_dy_idx = 37, color = [255, 0, 170]) + config.add_connection(from_body_part = BodyPart.right_shoulder, to_body_part = BodyPart.right_ear, paf_dx_idx = 18, paf_dy_idx = 19, color = [255, 0, 85]) + config.add_connection(from_body_part = BodyPart.left_shoulder, to_body_part = BodyPart.left_ear, paf_dx_idx = 26, paf_dy_idx = 27, color = [255, 0, 85]) + + return config \ No newline at end of file diff --git a/estimation_example/connections.py b/estimation_example/connections.py new file mode 100644 index 0000000..6878a8c --- /dev/null +++ b/estimation_example/connections.py @@ -0,0 +1,234 @@ +import math +import numpy as np +from config import get_default_configuration + + +def get_connections( + config, coords, paf, threshold=0.05, mid_num=10, minimum_mid_num=8): + """ + Finds the connection candidates and returns only valid connections. + + :param config: pose estimation configuration. + :param coords: dictionary with coordinates of all body parts. + :param paf: paf maps. + :param threshold: threshold for the intensity value in paf for a given mid point. If value at a mid point + is below the threshold the mid point is not taken into account. + :param mid_num: number of mid point for sampling + :param minimum_mid_num: minimum number of valid mid points for the connection candidate + :return: list of arrays containing identified connections of a given type : + [ + array( + [id1, id2, score1, score2, total_score] + [id1, id2, score1, score2, total_score] + ... + ), + array( + ... + ) + ] + """ + all_cand_connections = [] + + for conn in config.connection_types: + + # select dx and dy PAFs for this connection type + paf_dx = paf[:, :, conn.paf_dx_idx] + paf_dy = paf[:, :, conn.paf_dy_idx] + + # get coordinates lists for 2 body part types which belong to the current connection type + cand_a = coords[conn.from_body_part.name] + cand_b = coords[conn.to_body_part.name] + + n_a = len(cand_a) + n_b = len(cand_b) + max_connections = min(n_a, n_b) + + # lets check each combination of detected 2 body parts - candidate connections + if n_a != 0 and n_b != 0: + + # here we will store the connection candidates 5 columns: + # [ body part id1, body part id2, body part score1, body part score2, total score of connection ] + connection_candidates = np.zeros((0, 5)) + + for i in range(n_a): + for j in range(n_b): + # find the distance between the 2 body parts. The expression cand_b[j][:2] + # returns an 2 element array with coordinates x,y + vec = np.subtract(cand_b[j][:2], cand_a[i][:2]) + norm = math.sqrt(vec[0] * vec[0] + vec[1] * vec[1]) + + # skip the connection if 2 body parts overlaps + if norm == 0: + continue + + # normalize the vector + vec = np.divide(vec, norm) + + # get the set midpoints between 2 body parts (their coordinates x,y) + start_end = list(zip(np.linspace(cand_a[i][0], cand_b[j][0], num=mid_num), + np.linspace(cand_a[i][1], cand_b[j][1], num=mid_num))) + + # having the coord of midpoint we can read the intensity value in paf map at the midpoint + # for dx component + vec_x = np.array( + [paf_dx[int(round(start_end[i][1])), int( + round(start_end[i][0]))] for i in range(mid_num)] + ) + # for dy component + vec_y = np.array( + [paf_dy[int(round(start_end[i][1])), int( + round(start_end[i][0]))] for i in range(mid_num)] + ) + + # calculate the score for the connection weighted by the distance between body parts + score_midpts = np.multiply( + vec_x, vec[0]) + np.multiply(vec_y, vec[1]) + + # get the total score + total_score = sum(score_midpts) / len(score_midpts) + + # number of midpoints with intensity above the threshold shouldn't be less than 80% of all midpoints + criterion1 = len(np.nonzero( + score_midpts > threshold)[0]) > minimum_mid_num + criterion2 = total_score > 0 + + if criterion1 and criterion2: + # add this connection to the list [id1, id2, score1, score2, total score] + connection_candidates = np.vstack( + [connection_candidates, + [cand_a[i][3], + cand_b[j][3], + cand_a[i][2], + cand_b[j][2], + total_score]]) + + # sort the array by the total score - descending. (the sorted array is reversed by the expression [::-1]) + sorted_connections = connection_candidates[ + connection_candidates[:, 4].argsort()][::-1] + + # make sure we get no more than max_connections + all_cand_connections.append( + sorted_connections[:max_connections, :]) + else: + # not found any body parts but we still need to add empty list to preserve the correct indexing in the + # output array + all_cand_connections.append([]) + + return all_cand_connections + + +if __name__ == '__main__': + + coords = {'nose': + [(173, 13, 0.92409194, 0), + (85, 23, 0.9313662, 1), + (135, 29, 0.9052348, 2), + (19, 79, 0.9306832, 3), + (48, 83, 0.9516923, 4)], + 'neck': + [(172, 24, 0.8865844, 5), + (85, 33, 0.91056985, 6), + (129, 42, 0.7325343, 7), + (18, 89, 0.8726025, 8), + (47, 90, 0.9188747, 9)], + 'right_shoulder': + [(164, 26, 0.8121046, 10), + (76, 34, 0.88929117, 11), + (117, 42, 0.6240694, 12), + (11, 89, 0.85033226, 13), + (39, 90, 0.8911759, 14)], + 'right_elbow': + [(153, 39, 0.91216505, 15), + (97, 42, 0.39170936, 16), + (73, 49, 0.7160349, 17), + (18, 101, 0.617831, 18), + (38, 106, 0.78751093, 19)], + 'right_wrist': + [(160, 51, 0.6846192, 20), + (80, 52, 0.65054333, 21), + (24, 103, 0.26212907, 22), + (47, 111, 0.771493, 23)], + 'left_shoulder': + [(181, 22, 0.8704748, 24), + (95, 32, 0.85670036, 25), + (141, 41, 0.6815984, 26), + (26, 89, 0.82121867, 27), + (56, 90, 0.88108903, 28)], + 'left_elbow': + [(185, 37, 0.81237817, 29), + (99, 44, 0.22556348, 30), + (152, 56, 0.4836958, 31), + (34, 98, 0.67423993, 32), + (57, 105, 0.72111076, 33)], + 'left_wrist': + [(185, 49, 0.71519095, 34), + (158, 57, 0.44307223, 35), + (27, 103, 0.6169104, 36), + (52, 111, 0.5472575, 37)], + 'right_hip': + [(172, 59, 0.64131504, 38), + (82, 63, 0.62906384, 39), + (122, 96, 0.3766609, 40), + (43, 112, 0.6915504, 41), + (10, 115, 0.6116446, 42)], + 'right_knee': + [(87, 83, 0.6413851, 43), + (174, 83, 0.833872, 44), + (22, 101, 0.5693337, 45), + (37, 110, 0.34187955, 46), + (142, 128, 0.52756023, 47)], + 'right_ankle': + [(173, 106, 0.646768, 48), + (90, 108, 0.621414, 49), + (32, 119, 0.54047376, 50), + (53, 119, 0.29182506, 51), + (149, 167, 0.4758099, 52)], + 'left_hip': + [(182, 57, 0.6341558, 53), + (95, 62, 0.6280589, 54), + (129, 94, 0.3196735, 55), + (53, 111, 0.7400102, 56), + (22, 114, 0.46316183, 57)], + 'left_knee': + [(96, 82, 0.74924064, 58), + (182, 82, 0.83847153, 59), + (31, 100, 0.41720843, 60), + (67, 108, 0.85580474, 61), + (130, 136, 0.694633, 62)], + 'left_ankle': + [(83, 102, 0.658298, 63), + (179, 104, 0.72656405, 64), + (39, 120, 0.4767233, 65), + (46, 121, 0.57901657, 66), + (115, 166, 0.49369463, 67)], + 'right_eye': + [(170, 11, 0.9424342, 68), + (83, 21, 0.9409762, 69), + (131, 25, 0.8993071, 70), + (17, 77, 0.9191763, 71), + (46, 81, 0.9394402, 72)], + 'left_eye': + [(174, 11, 0.936771, 73), + (86, 21, 0.93059033, 74), + (137, 25, 0.92142344, 75), + (20, 78, 0.9070393, 76), + (50, 81, 0.9689289, 77)], + 'right_ear': + [(167, 13, 0.92250913, 78), + (80, 23, 0.90719926, 79), + (123, 27, 0.8572976, 80), + (13, 79, 0.8653215, 81), + (44, 81, 0.82936436, 82)], + 'left_ear': + [(177, 11, 0.77801484, 83), + (89, 22, 0.83940786, 84), + (23, 79, 0.81429726, 85), + (53, 82, 0.92744666, 86)]} + + paf_path = './resources/pafs.npy' + paf = np.load(paf_path) + + cfg = get_default_configuration() + connections = get_connections(cfg, coords, paf) + + print(connections) diff --git a/estimation_example/coordinates.py b/estimation_example/coordinates.py new file mode 100644 index 0000000..45f1ad5 --- /dev/null +++ b/estimation_example/coordinates.py @@ -0,0 +1,65 @@ +import numpy as np +from scipy.ndimage.filters import gaussian_filter +from config import get_default_configuration + + +def get_coordinates(config, heatmaps, threshold=0.1): + """ + Finds the coordinates x,y in the heatmaps. + + :param config: pose estimation configuration + :param heatmaps: heatmaps + :param threshold: threshold for the intensity value in the heatmap at the position of a peak + :return: dictionary: + { body_part_name: + [(x ,y, score, id), + (x ,y, score, id), + ... + ], + ... + } + """ + all_peaks = dict() + peak_counter = 0 + + for part_meta in config.body_parts.values(): + hmap_orig = heatmaps[:, :, part_meta.heatmap_idx] + hmap = gaussian_filter(hmap_orig, sigma=3) + + hmap_right = np.zeros(hmap.shape) + hmap_right[:, 1:] = hmap[:, :-1] + hmap_left = np.zeros(hmap.shape) + hmap_left[:, :-1] = hmap[:, 1:] + hmap_down = np.zeros(hmap.shape) + hmap_down[1:, :] = hmap[:-1, :] + hmap_up = np.zeros(hmap.shape) + hmap_up[:-1, :] = hmap[1:, :] + + peaks_binary = np.logical_and.reduce( + (hmap >= hmap_right, + hmap >= hmap_left, + hmap >= hmap_down, + hmap >= hmap_up, + hmap > threshold)) + peaks = list(zip(np.nonzero(peaks_binary)[1], + np.nonzero(peaks_binary)[0])) + + sequence_numbers = range(peak_counter, peak_counter + len(peaks)) + peaks_with_score_and_id = [peak + ( + hmap_orig[peak[1], peak[0]], + seq_num) for peak, seq_num in zip(peaks, sequence_numbers)] + + all_peaks[part_meta.body_part.name] = peaks_with_score_and_id + peak_counter += len(peaks) + + return all_peaks + + +if __name__ == '__main__': + heatmaps_path = './resources/heatmaps.npy' + heatmaps = np.load(heatmaps_path) + + cfg = get_default_configuration() + coordinates = get_coordinates(cfg, heatmaps) + + print(coordinates) diff --git a/estimation_example/estimators.py b/estimation_example/estimators.py new file mode 100644 index 0000000..a64c3e7 --- /dev/null +++ b/estimation_example/estimators.py @@ -0,0 +1,209 @@ +import math +import cv2 +import numpy as np +from config import get_default_configuration + + +def estimate(config, connections, min_num_body_parts=4, min_score=0.4): + """ + Estimates the skeletons. + + :param config: pose estimation configuration. + :param connections: valid connections + :param min_num_body_parts: minimum number of body parts for a skeleton + :param min_score: minimum score value for the skeleton + :return: list of skeletons. Each skeleton has a list of identifiers of body parts: + [ + [id1, id2,...,idN, parts_num, score], + [id1, id2,...,idN, parts_num, score] + ... + ] + """ + + # 2 extra slots for number of valid parts and overall score + number_of_slots = config.body_parts_size() + 2 + + # the connections are solely used to group body parts into separate skeletons. As a result we will + # get an array where each row represents a skeleton, each column contains an identifier of + # specific body part belonging to a skeleton (plus 2 extra columns for: num of parts and skeleton total score) + # we will be adding the skeletons to this array: + subset = np.empty((0, number_of_slots)) + + for k, conn in enumerate(config.connection_types): + if len(connections[k]) > 0: + # retrieve id and score of all body parts pairs for the current connection type + part_a = connections[k][:, [0, 2]] # idA, scoreA + part_b = connections[k][:, [1, 2]] # idB, scoreB + + # determine the slot number for 2 body parts types + slot_idx_a = config.body_parts[conn.from_body_part].slot_idx + slot_idx_b = config.body_parts[conn.to_body_part].slot_idx + + # iterate over all connection candidates filling up the subset with the correct body part identifiers + for i in range(len(connections[k])): + found = 0 + slot_idx = [-1, -1] + for j in range(len(subset)): + if subset[j][slot_idx_a] == part_a[i, 0] or \ + subset[j][slot_idx_b] == part_b[i, 0]: + slot_idx[found] = j + found += 1 + + if found == 1: + j = slot_idx[0] + if subset[j][slot_idx_b] != part_b[i, 0]: + subset[j][slot_idx_b] = part_b[i, 0] + subset[j][-1] += 1 + subset[j][-2] += part_b[i, 1] + connections[k][i][2] + + elif found == 2: # if found 2 and disjoint, merge them + j1, j2 = slot_idx + membership = ((subset[j1] >= 0).astype(int) + + (subset[j2] >= 0).astype(int))[:-2] + if len(np.nonzero(membership == 2)[0]) == 0: # merge + subset[j1][:-2] += (subset[j2][:-2] + 1) + subset[j1][-2:] += subset[j2][-2:] + subset[j1][-2] += connections[k][i][2] + subset = np.delete(subset, j2, 0) + else: # as like found == 1 + subset[j1][slot_idx_b] = part_b[i, 0] + subset[j1][-1] += 1 + subset[j1][-2] += part_b[i, 1] + connections[k][i][2] + + # if find no partA in the subset, create a new subset + elif not found: + row = -1 * np.ones(number_of_slots) + row[slot_idx_a] = part_a[i, 0] + row[slot_idx_b] = part_b[i, 0] + row[-1] = 2 + row[-2] = part_a[i, 1] + part_b[i, 1] + connections[k][i][2] + subset = np.vstack([subset, row]) + + # delete some rows of subset which has few parts occur + delete_idx = [] + for i in range(len(subset)): + if subset[i][-1] < min_num_body_parts or \ + subset[i][-2] / subset[i][-1] < min_score: + delete_idx.append(i) + subset = np.delete(subset, delete_idx, axis=0) + + return subset + + +if __name__ == '__main__': + connections = [ + np.array( + [[ 9. , 14. , 0.9188747 , 0.8911759 , 1.03683093], + [ 5. , 10. , 0.8865844 , 0.8121046 , 0.94687685], + [ 8. , 13. , 0.8726025 , 0.85033226, 0.9012778 ], + [ 6. , 11. , 0.91056985, 0.88929117, 0.88995027], + [ 7. , 12. , 0.7325343 , 0.6240694 , 0.78442425]]), + np.array( + [[ 9. , 28. , 0.9188747 , 0.88108903, 1.00724218], + [ 8. , 27. , 0.8726025 , 0.82121867, 0.96535013], + [ 5. , 24. , 0.8865844 , 0.8704748 , 0.93374986], + [ 6. , 25. , 0.91056985, 0.85670036, 0.92527213], + [ 7. , 26. , 0.7325343 , 0.6815984 , 0.84348354]]), + np.array( + [[10. , 15. , 0.8121046 , 0.91216505, 0.99742639], + [14. , 19. , 0.8911759 , 0.78751093, 0.97054861], + [13. , 18. , 0.85033226, 0.617831 , 0.91969185], + [11. , 17. , 0.88929117, 0.7160349 , 0.83139618], + [12. , 16. , 0.6240694 , 0.39170936, 0.47590637]]), + np.array( + [[15. , 20. , 0.91216505, 0.6846192 , 1.09104793], + [19. , 23. , 0.78751093, 0.771493 , 0.89008226], + [18. , 22. , 0.617831 , 0.26212907, 0.63477109], + [18. , 23. , 0.617831 , 0.771493 , 0.58728189]]), + np.array( + [[28. , 33. , 0.88108903, 0.72111076, 1.00839311], + [27. , 32. , 0.82121867, 0.67423993, 0.98828157], + [24. , 29. , 0.8704748 , 0.81237817, 0.9666315 ], + [26. , 31. , 0.6815984 , 0.4836958 , 0.80339349], + [25. , 30. , 0.85670036, 0.22556348, 0.51399985]]), + np.array( + [[29. , 34. , 0.81237817, 0.71519095, 0.90577437], + [33. , 37. , 0.72111076, 0.5472575 , 0.65625855], + [32. , 36. , 0.67423993, 0.6169104 , 0.62864926], + [31. , 35. , 0.4836958 , 0.44307223, 0.60174478]]), + np.array( + [[ 9. , 41. , 0.9188747 , 0.6915504 , 0.9004276 ], + [ 5. , 38. , 0.8865844 , 0.64131504, 0.85271515], + [ 8. , 42. , 0.8726025 , 0.6116446 , 0.83352115], + [ 6. , 39. , 0.91056985, 0.62906384, 0.83043894], + [ 7. , 40. , 0.7325343 , 0.3766609 , 0.62170903]]), + np.array( + [[38. , 44. , 0.64131504, 0.833872 , 0.85671546], + [39. , 43. , 0.62906384, 0.6413851 , 0.85363647], + [42. , 45. , 0.6116446 , 0.5693337 , 0.70891048], + [40. , 47. , 0.3766609 , 0.52756023, 0.66385489], + [41. , 46. , 0.6915504 , 0.34187955, 0.4365968 ]]), + np.array( + [[44. , 48. , 0.833872 , 0.646768 , 0.9471579 ], + [45. , 50. , 0.5693337 , 0.54047376, 0.90375594], + [43. , 49. , 0.6413851 , 0.621414 , 0.88494388], + [47. , 52. , 0.52756023, 0.4758099 , 0.71417868], + [46. , 51. , 0.34187955, 0.29182506, 0.57783459]]), + np.array( + [[ 6. , 54. , 0.91056985, 0.6280589 , 0.96156364], + [ 9. , 56. , 0.9188747 , 0.7400102 , 0.960601 ], + [ 5. , 53. , 0.8865844 , 0.6341558 , 0.94478426], + [ 8. , 57. , 0.8726025 , 0.46316183, 0.87909239], + [ 7. , 55. , 0.7325343 , 0.3196735 , 0.60074055]]), + np.array( + [[53. , 59. , 0.6341558 , 0.83847153, 0.84733025], + [54. , 58. , 0.6280589 , 0.74924064, 0.81335635], + [56. , 61. , 0.7400102 , 0.85580474, 0.81244228], + [55. , 62. , 0.3196735 , 0.694633 , 0.5182847 ], + [57. , 61. , 0.46316183, 0.85580474, 0.39467257]]), + np.array( + [[59. , 64. , 0.83847153, 0.72656405, 0.88302198], + [61. , 66. , 0.85580474, 0.57901657, 0.82494843], + [58. , 63. , 0.74924064, 0.658298 , 0.82203498], + [62. , 67. , 0.694633 , 0.49369463, 0.74128318], + [60. , 65. , 0.41720843, 0.4767233 , 0.453668 ]]), + np.array( + [[6. , 1. , 0.91056985, 0.9313662 , 0.99411402], + [8. , 3. , 0.8726025 , 0.9306832 , 0.98572515], + [9. , 4. , 0.9188747 , 0.9516923 , 0.98134359], + [5. , 0. , 0.8865844 , 0.92409194, 0.94417737], + [7. , 2. , 0.7325343 , 0.9052348 , 0.84682636]]), + np.array( + [[ 0. , 68. , 0.92409194, 0.9424342 , 1.14795679], + [ 3. , 71. , 0.9306832 , 0.9191763 , 1.130467 ], + [ 1. , 69. , 0.9313662 , 0.9409762 , 1.10818533], + [ 2. , 70. , 0.9052348 , 0.8993071 , 1.10167558], + [ 4. , 72. , 0.9516923 , 0.9394402 , 1.08395883]]), + np.array( + [[72. , 82. , 0.9394402 , 0.82936436, 0.9980747 ], + [68. , 78. , 0.9424342 , 0.92250913, 0.98662932], + [69. , 79. , 0.9409762 , 0.90719926, 0.98011624], + [70. , 80. , 0.8993071 , 0.8572976 , 0.87976664], + [71. , 81. , 0.9191763 , 0.8653215 , 0.7991888 ]]), + np.array( + [[ 2. , 75. , 0.9052348 , 0.92142344, 0.99353639], + [ 0. , 73. , 0.92409194, 0.936771 , 0.9643986 ], + [ 3. , 76. , 0.9306832 , 0.9070393 , 0.95751885], + [ 4. , 77. , 0.9516923 , 0.9689289 , 0.95725774], + [ 1. , 74. , 0.9313662 , 0.93059033, 0.95717775]]), + np.array( + [[77. , 86. , 0.9689289 , 0.92744666, 1.03445414], + [74. , 84. , 0.93059033, 0.83940786, 0.99919387], + [76. , 85. , 0.9070393 , 0.81429726, 0.95994219], + [73. , 83. , 0.936771 , 0.77801484, 0.87463949]]), + np.array( + [[14. , 82. , 0.8911759 , 0.82936436, 0.95386559], + [10. , 78. , 0.8121046 , 0.92250913, 0.911006 ], + [11. , 79. , 0.88929117, 0.90719926, 0.89117094], + [13. , 81. , 0.85033226, 0.8653215 , 0.84301486], + [12. , 80. , 0.6240694 , 0.8572976 , 0.80546115]]), + np.array( + [[28. , 86. , 0.88108903, 0.92744666, 1.04618902], + [25. , 84. , 0.85670036, 0.83940786, 0.98009169], + [24. , 83. , 0.8704748 , 0.77801484, 0.93649059], + [27. , 85. , 0.82121867, 0.81429726, 0.929685 ]])] + + cfg = get_default_configuration() + skeletons = estimate(cfg, connections) + + print(skeletons) diff --git a/estimation_example/example.py b/estimation_example/example.py new file mode 100644 index 0000000..348eeda --- /dev/null +++ b/estimation_example/example.py @@ -0,0 +1,34 @@ +import cv2 +import numpy as np + +from config import get_default_configuration +from coordinates import get_coordinates +from connections import get_connections +from estimators import estimate +from renderers import draw + + +if __name__ == '__main__': + heatmaps_path = './resources/heatmaps.npy' + paf_path = './resources/pafs.npy' + example_img_path = 'resources/ski.jpg' + output_img_path = 'output.jpg' + + example_image = cv2.imread(example_img_path) + + heatmaps = np.load(heatmaps_path) + paf = np.load(paf_path) + + cfg = get_default_configuration() + + coordinates = get_coordinates(cfg, heatmaps) + + connections = get_connections(cfg, coordinates, paf) + + skeletons = estimate(cfg, connections) + + output = draw(cfg, example_image, coordinates, skeletons) + + cv2.imwrite(output_img_path, output) + + print(f"Output image: {output_img_path}") diff --git a/estimation_example/output.jpg b/estimation_example/output.jpg new file mode 100644 index 0000000..faebde7 Binary files /dev/null and b/estimation_example/output.jpg differ diff --git a/estimation_example/renderers.py b/estimation_example/renderers.py new file mode 100644 index 0000000..cd575a3 --- /dev/null +++ b/estimation_example/renderers.py @@ -0,0 +1,50 @@ +import cv2 +import math +import numpy as np + +def draw(config, input_image, coords, subset, resize_fac = 1): + + stickwidth = 1 + + canvas = input_image.copy() + + for body_part_type, body_part_meta in config.body_parts.items(): + color = body_part_meta.color + body_part_peaks = coords[body_part_type.name] + + for peak in body_part_peaks: + a = peak[0] * resize_fac + b = peak[1] * resize_fac + cv2.circle(canvas, (a, b), stickwidth, color, thickness=-1) + + # dict(id: [y,x]) Note, coord are reversed + xy_by_id = dict([(item[3], np.array([item[1], item[0]])) for sublist in coords.values() for item in sublist]) + + xy = np.zeros((2,2)) + for i, conn_type in enumerate(config.connection_types): + index1 = config.body_parts[conn_type.from_body_part].slot_idx + index2 = config.body_parts[conn_type.to_body_part].slot_idx + indexes = np.array([index1, index2]) + for s in subset: + + ids = s[indexes] + if -1 in ids: + continue + + cur_canvas = canvas.copy() + xy[0, :] = xy_by_id[ids[0]] + xy[1, :] = xy_by_id[ids[1]] + + m_x = np.mean(xy[:, 0]) + m_y = np.mean(xy[:, 1]) + sss = xy[1, 1] + length = ((xy[0, 0] - xy[1, 0]) ** 2 + (xy[0, 1] - xy[1, 1]) ** 2) ** 0.5 + + angle = math.degrees(math.atan2(xy[0, 0] - xy[1, 0], xy[0, 1] - xy[1, 1])) + + polygon = cv2.ellipse2Poly((int(m_y * resize_fac), int(m_x * resize_fac)), + (int(length * resize_fac / 2), stickwidth), int(angle), 0, 360, 1) + cv2.fillConvexPoly(cur_canvas, polygon, conn_type.color) + canvas = cv2.addWeighted(canvas, 0.4, cur_canvas, 0.6, 0) + + return canvas \ No newline at end of file diff --git a/estimation_example/resources/heatmaps.npy b/estimation_example/resources/heatmaps.npy new file mode 100644 index 0000000..2957eae Binary files /dev/null and b/estimation_example/resources/heatmaps.npy differ diff --git a/estimation_example/resources/pafs.npy b/estimation_example/resources/pafs.npy new file mode 100644 index 0000000..eed695a Binary files /dev/null and b/estimation_example/resources/pafs.npy differ diff --git a/estimation_example/resources/ski.jpg b/estimation_example/resources/ski.jpg new file mode 100755 index 0000000..5a52e4a Binary files /dev/null and b/estimation_example/resources/ski.jpg differ