#!/usr/bin/env python # coding: utf-8 # # OpenPose pose detection # 2D real-time multi-person keypoint detection: # **18**-keypoint body/foot keypoint estimation. Running time invariant to number of detected people # see https://github.com/CMU-Perceptual-Computing-Lab/openpose # # ## Pipeline # - Run 18-keypoint model on video frames # - Parse keypoints and PAFs to generate personwise keypoints # - Save results to OpenPose.pkl import os import numpy as np import cv2 import sys from sys import platform import time import pandas as pd import matplotlib.pyplot as plt from tqdm import tqdm, tqdm_pandas tqdm.pandas() from multiprocessing import cpu_count from multiprocessing import Pool import itertools import os os.environ["CUDA_VISIBLE_DEVICES"] = "1" os.environ["OPENCV_DNN_OPENCL_ALLOW_ALL_DEVICES"] = "1" POSE_PAIRS = [[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, 17], [5, 16]] keypointsMapping = ['Nose', 'Neck', 'R-Sho', 'R-Elb', 'R-Wr', 'L-Sho', 'L-Elb', 'L-Wr', 'R-Hip', 'R-Knee', 'R-Ank', 'L-Hip', 'L-Knee', 'L-Ank', 'R-Eye', 'L-Eye', 'R-Ear', 'L-Ear'] mapIdx = [[31, 32], [39, 40], [33, 34], [35, 36], [41, 42], [43, 44], [19, 20], [21, 22], [23, 24], [25, 26], [27, 28], [29, 30], [47, 48], [49, 50], [53, 54], [51, 52], [55, 56], [37, 38], [45, 46]] colors = np.array([(0, 100, 255), (0, 100, 255), (0, 255, 255), (0, 100, 255), (0, 255, 255), (0, 100, 255), (0, 255, 0), (255, 200, 100), (255, 0, 255), (0, 255, 0), (255, 200, 100), (255, 0, 255), (0, 0, 255), (255, 0, 0), (200, 200, 0), (255, 0, 0), (200, 200, 0), (0, 0, 0)]) my_color = [] for c in colors: my_color.append(tuple(c)) # ## Auxiliary Functions # see https://www.learnopencv.com/deep-learning-based-human-pose-estimation-using-opencv-cpp-python/ def getKeypoints(probMap, threshold=0.8): mapSmooth = cv2.GaussianBlur(probMap, (3, 3), 0, 0) mapMask = np.uint8(mapSmooth > threshold) keypoints = [] # find the blobs contours, _ = cv2.findContours(mapMask, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE) # for each blob find the maxima # not enough values to unpack (expected 3, got 2) # version issue: https://github.com/facebookresearch/maskrcnn-benchmark/issues/339 for cnt in contours: blobMask = np.zeros(mapMask.shape) blobMask = cv2.fillConvexPoly(blobMask, cnt, 1) maskedProbMap = mapSmooth * blobMask _, maxVal, _, maxLoc = cv2.minMaxLoc(maskedProbMap) keypoints.append(maxLoc + (probMap[maxLoc[1], maxLoc[0]],)) return keypoints # Find valid connections between the different joints of a all persons present def getValidPairs(output, detected_keypoints, frameWidth, frameHeight): valid_pairs = [] invalid_pairs = [] n_interp_samples = 10 paf_score_th = 0.1 conf_th = 0.7 # loop for every POSE_PAIR for k in range(len(mapIdx)): # A->B constitute a limb pafA = output[mapIdx[k][0], :, :] pafB = output[mapIdx[k][1], :, :] pafA = cv2.resize(pafA, (frameWidth, frameHeight)) pafB = cv2.resize(pafB, (frameWidth, frameHeight)) # Find the keypoints for the first and second limb candA = detected_keypoints[POSE_PAIRS[k][0]] candB = detected_keypoints[POSE_PAIRS[k][1]] nA = len(candA) nB = len(candB) # If keypoints for the joint-pair is detected # check every joint in candA with every joint in candB # Calculate the distance vector between the two joints # Find the PAF values at a set of interpolated points between the joints # Use the above formula to compute a score to mark the connection valid if (nA != 0 and nB != 0): valid_pair = np.zeros((0, 3)) for i in range(nA): max_j = -1 maxScore = -1 found = 0 for j in range(nB): # Find d_ij d_ij = np.subtract(candB[j][:2], candA[i][:2]) norm = np.linalg.norm(d_ij) if norm: d_ij = d_ij / norm else: continue # Find p(u) interp_coord = list(zip(np.linspace(candA[i][0], candB[j][0], num=n_interp_samples), np.linspace(candA[i][1], candB[j][1], num=n_interp_samples))) # Find L(p(u)) paf_interp = [] for k in range(len(interp_coord)): paf_interp.append([pafA[int(round(interp_coord[k][1])), int(round(interp_coord[k][0]))], pafB[int(round(interp_coord[k][1])), int(round(interp_coord[k][0]))]]) # Find E paf_scores = np.dot(paf_interp, d_ij) avg_paf_score = sum(paf_scores) / len(paf_scores) # Check if the connection is valid # If the fraction of interpolated vectors aligned with PAF is higher then threshold -> Valid Pair if (len(np.where(paf_scores > paf_score_th)[0]) / n_interp_samples) > conf_th: if avg_paf_score > maxScore: max_j = j maxScore = avg_paf_score found = 1 # Append the connection to the list if found: valid_pair = np.append(valid_pair, [[candA[i][3], candB[max_j][3], maxScore]], axis=0) # Append the detected connections to the global list valid_pairs.append(valid_pair) else: # If no keypoints are detected # ATTENTION: Commented this out by Sven # print("No Connection : k = {}".format(k)) invalid_pairs.append(k) valid_pairs.append([]) return valid_pairs, invalid_pairs # This function creates a list of keypoints belonging to each person # For each detected valid pair, it assigns the joint(s) to a person def getPersonwiseKeypoints(valid_pairs, invalid_pairs, keypoints_list): # the last number in each row is the overall score personwiseKeypoints = -1 * np.ones((0, 19)) for k in range(len(mapIdx)): if k not in invalid_pairs: partAs = valid_pairs[k][:, 0] partBs = valid_pairs[k][:, 1] indexA, indexB = np.array(POSE_PAIRS[k]) for i in range(len(valid_pairs[k])): found = 0 person_idx = -1 for j in range(len(personwiseKeypoints)): if personwiseKeypoints[j][indexA] == partAs[i]: person_idx = j found = 1 break if found: personwiseKeypoints[person_idx][indexB] = partBs[i] personwiseKeypoints[person_idx][-1] += keypoints_list[partBs[i].astype(int), 2] + valid_pairs[k][i][ 2] # if find no partA in the subset, create a new subset elif not found and k < 17: row = -1 * np.ones(19) row[indexA] = partAs[i] row[indexB] = partBs[i] # add the keypoint_scores for the two keypoints and the paf_score row[-1] = sum(keypoints_list[valid_pairs[k][i, :2].astype(int), 2]) + valid_pairs[k][i][2] personwiseKeypoints = np.vstack([personwiseKeypoints, row]) return personwiseKeypoints def f(probMap): threshold = 0.5 return getKeypoints(probMap, threshold) def getPose(output): detected_keypoints = [] keypoints_list = np.zeros((0, 3)) keypoint_id = 0 threshold = 0.5 keypointsList = [] for part in range(18): probMap = output[part, :, :] probMap = cv2.resize(probMap, (frameWidth, frameHeight)) keypointsList.append(getKeypoints(probMap, threshold)) for keypoints in keypointsList: # nPoints = 18 keypoints_with_id = [] for i in range(len(keypoints)): keypoints_with_id.append(keypoints[i] + (keypoint_id,)) keypoints_list = np.vstack([keypoints_list, keypoints[i]]) keypoint_id += 1 detected_keypoints.append(keypoints_with_id) valid_pairs, invalid_pairs = getValidPairs(output, detected_keypoints, frameWidth, frameHeight) personwiseKeypoints = getPersonwiseKeypoints(valid_pairs, invalid_pairs, keypoints_list) return detected_keypoints, keypoints_list, personwiseKeypoints """Forward array of 20 images""" def getPoseFromDNN(net, images, frameWidth, frameHeight): inHeight = 368 inWidth = int((inHeight / frameHeight) * frameWidth) inpBlob = cv2.dnn.blobFromImages(np.array(images), 1.0 / 255, (inWidth, inHeight), (0, 0, 0), swapRB=False, crop=False) # Set the prepared object as the input blob of the network net.setInput(inpBlob) output = net.forward() return output def visualize(image, df, frame): number_ids = len([col for col in df.columns if 'ID' in col]) data = df.loc[df['Frame'] == frame] plt.imshow(image) plt.axis('off') for id_no in range(number_ids): keypoints = df['ID%i_Keypoints' % id_no].iloc[frame] for i in range(len(POSE_PAIRS)): index = POSE_PAIRS[i] A, B = keypoints[index] # for idx in index: # print(keypointsMapping[idx]) if A is not None and B is not None: plt.plot((A[0], B[0]), (A[1], B[1]), c=colors[i]) plt.show() def process(file): global get_keypoints POSE_PROTO_FILE = r"openpose/pose_deploy_linevec.prototxt" POSE_WEIGHTS_FILE = r"openpose/pose_iter_440000.caffemodel" if not os.path.exists(POSE_PROTO_FILE): print('WARNING: Could not find pose file %s' % POSE_PROTO_FILE) return if not os.path.exists(POSE_WEIGHTS_FILE): print('WARNING: Could not find model weights file %s' % POSE_WEIGHTS_FILE) return VIDEO = file ROOT = "/".join(VIDEO.split("/")[:-1]) + "/" VIDEOOUT = VIDEO.split("/")[-1].split(".")[0] TMP_DIR = "/".join(VIDEO.split("/")[:-2]) + "/temp/" FRAMES = "%s%s_frames" % (TMP_DIR, VIDEOOUT) if not os.path.exists(FRAMES): print('WARNING: Could not find frame directory') return # Load Model # net = cv2.dnn.readNetFromCaffe(POSE_PROTO_FILE, POSE_WEIGHTS_FILE) net.setPreferableBackend(cv2.dnn.DNN_BACKEND_CUDA) net.setPreferableTarget(cv2.dnn.DNN_TARGET_CUDA) path_list = [f for f in os.listdir(FRAMES) if '.jpg' in f] path_list.sort() image = cv2.imread(os.path.join(FRAMES, path_list[0])) image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) frameWidth = image.shape[1] frameHeight = image.shape[0] lst = [] images = [] for frame, path in enumerate(tqdm(path_list)): image = cv2.imread(os.path.join(FRAMES, path)) image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) images.append(image) if len(images) == 20: output = getPoseFromDNN(net, images, frameWidth, frameHeight) lst.extend(output) images = [] if len(images) != 0: output = getPoseFromDNN(net, images, frameWidth, frameHeight) lst.extend(output) images = [] print(len(lst)) df = pd.DataFrame(range(len(lst))) df.columns = ["Frame"] df["Pose"] = lst path = "%s%s_OpenPose_raw.pkl" % (TMP_DIR, VIDEOOUT) df.to_pickle(path) def get_keypoints(frames): ret = [] for f in frames: output = df[df.Frame == f].Pose.iloc[0] ret.append(getPose(output)) return ret if platform == "linux" or platform == "linux2" or platform == "darwin": cores = max(4, cpu_count() - 4) print('Run sorting on {} cores'.format(cores)) data_split = np.array_split(df.Frame, cores) pool = Pool(cores) data = pool.map(get_keypoints, data_split) pool.close() pool.join() elif platform == "win32": print('WARNING: Can\'t run multiprocessing on Windows, this might take a while...') data = get_keypoints(df.Frame) else: print('WARNING: platform not supported') x = np.vstack(data) df["DetectedKeypoints"] = x[:, 0] df["KeypointsList"] = x[:, 1] df["PersonwiseKeypoints"] = x[:, 2] df.head() path = "%s%s_OpenPose.pkl" % (TMP_DIR, VIDEOOUT) df.to_pickle(path) print("Saved OpenPose detections to %s" % path) del df["Pose"] del df['Pic'] del df['DetectedKeypoints'] number_ids = len(df.PersonwiseKeypoints.values.flatten()[0]) print('Number of detected IDs: ', number_ids) """Map personwise keypoints to list of keypoints for each ID""" def keypoints_fun(x): # Discard frames where not all ids detected if len(x.PersonwiseKeypoints) < number_ids: # print('None') return None # index is -1 for no detection >> keypoint = None lst = list(x.KeypointsList) lst.append(None) lst = np.array(lst) keypoints = lst[x.PersonwiseKeypoints[id_no].astype(int)[:18]] return keypoints for id_no in range(number_ids): counter = 0 print('ID%i' % id_no) col = 'ID%i_Keypoints' % id_no df[col] = df.apply(keypoints_fun, axis=1) """Sort IDs to be consistent throughout video""" lst = [] columns = [col for col in df.columns if 'ID' in col] data = df[columns] lst.append(data.iloc[0].values) for i in range(1, len(df.Frame)): row = data.iloc[i] lst2 = [] for ids in range(number_ids): keypoints = row['ID%i_Keypoints' % ids] if keypoints is not None and keypoints[1] is not None: for j in range(number_ids): backtrack = 1 while lst[i - backtrack][j] == None: backtrack = backtrack + 1 keypoints2 = lst[i - backtrack][j] lst2.append([ids, j, np.linalg.norm(np.array(keypoints[1]) - np.array(keypoints2[1])), keypoints]) else: lst2.append([ids, None, None, None]) dfX = pd.DataFrame(lst2) dfX.columns = ["Id", "GtId", "Distance", "Keypoints"] dfX = dfX.sort_values("Distance") dfX = dfX.drop_duplicates("GtId").drop_duplicates("Id") lstRow = [] for j in range(number_ids): if (len(dfX[dfX.GtId == j]) > 0): lstRow.append(dfX[dfX.GtId == j].iloc[0].Keypoints) else: lstRow.append(None) lstRow.append(i) lst.append(lstRow) df_new = pd.DataFrame(lst) columns = [] for i in range(number_ids): columns.append('ID%i_Keypoints' % i) columns.append("Frame") df_new.columns = columns # First frame number is NaN from sorting df_new.Frame = df_new.Frame.fillna(0) df_new = df_new.astype({'Frame': 'int32'}) path = "%s%s_BodyMovement.pkl" % (TMP_DIR, VIDEOOUT) df_new.to_pickle(path) print('Saved Body Movement to %s' % path) visualize(image, df_new, 0)