From 6cce25792c9d950b39f75638fa4b3cfe2d3b386a Mon Sep 17 00:00:00 2001 From: shuaiqing Date: Mon, 28 Jun 2021 10:38:36 +0800 Subject: [PATCH] [mvmp] update track and fit smpl --- apps/demo/auto_track.py | 29 +++ apps/demo/mv1p_mirror.py | 4 +- apps/demo/smpl_from_keypoints.py | 62 ++++++ easymocap/affinity/matchSVT.py | 13 +- easymocap/assignment/track.py | 345 ++++++++++++++++++++++++++++++ easymocap/mytools/reader.py | 4 +- easymocap/mytools/utils.py | 6 +- easymocap/pipeline/weight.py | 26 ++- scripts/postprocess/eval_shelf.py | 289 +++++++++++++++++++++++++ 9 files changed, 760 insertions(+), 18 deletions(-) create mode 100644 apps/demo/auto_track.py create mode 100644 apps/demo/smpl_from_keypoints.py create mode 100644 easymocap/assignment/track.py create mode 100644 scripts/postprocess/eval_shelf.py diff --git a/apps/demo/auto_track.py b/apps/demo/auto_track.py new file mode 100644 index 0000000..78c3514 --- /dev/null +++ b/apps/demo/auto_track.py @@ -0,0 +1,29 @@ +''' + @ Date: 2021-06-25 15:59:35 + @ Author: Qing Shuai + @ LastEditors: Qing Shuai + @ LastEditTime: 2021-06-28 10:32:24 + @ FilePath: /EasyMocapRelease/apps/demo/auto_track.py +''' +from easymocap.assignment.track import Track2D, Track3D + +if __name__ == "__main__": + import argparse + parser = argparse.ArgumentParser() + parser.add_argument('path', type=str) + parser.add_argument('out', type=str) + parser.add_argument('--track3d', action='store_true') + parser.add_argument('--debug', action='store_true') + args = parser.parse_args() + cfg = { + 'path': args.path, + 'out': args.out, + 'WINDOW_SIZE': 10, + 'MIN_FRAMES': 10, + 'SMOOTH_SIZE': 5 + } + if args.track3d: + tracker = Track3D(with2d=False, **cfg) + else: + tracker = Track2D(**cfg) + tracker.auto_track() \ No newline at end of file diff --git a/apps/demo/mv1p_mirror.py b/apps/demo/mv1p_mirror.py index 8c858ef..18fc6d3 100644 --- a/apps/demo/mv1p_mirror.py +++ b/apps/demo/mv1p_mirror.py @@ -2,7 +2,7 @@ @ Date: 2021-04-13 22:21:39 @ Author: Qing Shuai @ LastEditors: Qing Shuai - @ LastEditTime: 2021-04-14 12:22:59 + @ LastEditTime: 2021-06-14 15:31:48 @ FilePath: /EasyMocap/apps/demo/mv1p_mirror.py ''' import os @@ -33,6 +33,6 @@ if __name__ == "__main__": if args.skel or not os.path.exists(skel_path): mv1pmf_skel(dataset, check_repro=False, args=args) from easymocap.pipeline.weight import load_weight_pose, load_weight_shape - weight_shape = load_weight_shape(args.opts) + weight_shape = load_weight_shape(args.model, args.opts) weight_pose = load_weight_pose(args.model, args.opts) mv1pmf_smpl(dataset, args=args, weight_pose=weight_pose, weight_shape=weight_shape) \ No newline at end of file diff --git a/apps/demo/smpl_from_keypoints.py b/apps/demo/smpl_from_keypoints.py new file mode 100644 index 0000000..0f75459 --- /dev/null +++ b/apps/demo/smpl_from_keypoints.py @@ -0,0 +1,62 @@ +''' + @ Date: 2021-06-14 22:27:05 + @ Author: Qing Shuai + @ LastEditors: Qing Shuai + @ LastEditTime: 2021-06-28 10:33:26 + @ FilePath: /EasyMocapRelease/apps/demo/smpl_from_keypoints.py +''' +# This is the script of fitting SMPL to 3d(+2d) keypoints +from easymocap.dataset import CONFIG +from easymocap.mytools import Timer +from easymocap.smplmodel import load_model, select_nf +from easymocap.mytools.reader import read_keypoints3d_all +from easymocap.mytools.file_utils import write_smpl +from easymocap.pipeline.weight import load_weight_pose, load_weight_shape +from easymocap.pipeline import smpl_from_keypoints3d +import os +from os.path import join +from tqdm import tqdm + +def smpl_from_skel(path, sub, out, skel3d, args): + config = CONFIG[args.body] + results3d, filenames = read_keypoints3d_all(skel3d) + pids = list(results3d.keys()) + weight_shape = load_weight_shape(args.model, args.opts) + weight_pose = load_weight_pose(args.model, args.opts) + with Timer('Loading {}, {}'.format(args.model, args.gender)): + body_model = load_model(args.gender, model_type=args.model) + for pid, result in results3d.items(): + body_params = smpl_from_keypoints3d(body_model, result['keypoints3d'], config, args, + weight_shape=weight_shape, weight_pose=weight_pose) + result['body_params'] = body_params + + # write for each frame + for nf, skelname in enumerate(tqdm(filenames, desc='writing')): + basename = os.path.basename(skelname) + outname = join(out, basename) + res = [] + for pid, result in results3d.items(): + frames = result['frames'] + if nf in frames: + nnf = frames.index(nf) + val = {'id': pid} + params = select_nf(result['body_params'], nnf) + val.update(params) + res.append(val) + write_smpl(outname, res) + +if __name__ == "__main__": + from easymocap.mytools import load_parser, parse_parser + parser = load_parser() + parser.add_argument('--skel3d', type=str, required=True) + args = parse_parser(parser) + help=""" + Demo code for fitting SMPL to 3d(+2d) skeletons: + + - Input : {} => {} + - Output: {} + - Body : {}=>{}, {} +""".format(args.path, args.skel3d, args.out, + args.model, args.gender, args.body) + print(help) + smpl_from_skel(args.path, args.sub, args.out, args.skel3d, args) \ No newline at end of file diff --git a/easymocap/affinity/matchSVT.py b/easymocap/affinity/matchSVT.py index caa9f9a..c51b0a4 100644 --- a/easymocap/affinity/matchSVT.py +++ b/easymocap/affinity/matchSVT.py @@ -2,19 +2,24 @@ @ Date: 2021-06-04 20:47:38 @ Author: Qing Shuai @ LastEditors: Qing Shuai - @ LastEditTime: 2021-06-04 21:50:53 - @ FilePath: /EasyMocapRelease/easymocap/affinity/matchSVT.py + @ LastEditTime: 2021-06-15 17:30:16 + @ FilePath: /EasyMocap/easymocap/affinity/matchSVT.py ''' import numpy as np -def matchSVT(M_aff, dimGroups, M_constr, M_obs, control): +def matchSVT(M_aff, dimGroups, M_constr=None, M_obs=None, control={}): max_iter = control['maxIter'] w_rank = control['w_rank'] tol = control['tol'] - X = M_aff + X = M_aff.copy() N = X.shape[0] index_diag = np.arange(N) X[index_diag, index_diag] = 0. + if M_constr is None: + M_constr = np.ones_like(M_aff) + for i in range(len(dimGroups) - 1): + M_constr[dimGroups[i]:dimGroups[i+1], dimGroups[i]:dimGroups[i+1]] = 0 + M_constr[index_diag, index_diag] = 1 X = (X + X.T)/2 Y = np.zeros((N, N)) mu = 64 diff --git a/easymocap/assignment/track.py b/easymocap/assignment/track.py new file mode 100644 index 0000000..3c18e95 --- /dev/null +++ b/easymocap/assignment/track.py @@ -0,0 +1,345 @@ +''' + @ Date: 2021-06-27 16:21:50 + @ Author: Qing Shuai + @ LastEditors: Qing Shuai + @ LastEditTime: 2021-06-28 10:32:40 + @ FilePath: /EasyMocapRelease/easymocap/assignment/track.py +''' +from tqdm import tqdm +import numpy as np +import os +from os.path import join +from glob import glob +from ..affinity.affinity import getDimGroups +from ..affinity.matchSVT import matchSVT +from ..mytools.reader import read_keypoints2d, read_keypoints3d +from ..mytools.file_utils import read_annot, read_json, save_annot, save_json, write_keypoints3d + +def check_path(x): + assert os.path.exists(x), '{} not exists!'.format(x) + +class BaseTrack: + def __init__(self, path, out, WINDOW_SIZE, MIN_FRAMES, SMOOTH_SIZE) -> None: + self.path = path + self.out = out + self.WINDOW_SIZE = WINDOW_SIZE + self.SMOOTH_SIZE = SMOOTH_SIZE + self.MIN_FRAMES = MIN_FRAMES + self.svt_args = { + 'maxIter': 1000, + 'w_sparse': 0.3, + 'w_rank': 10, + 'tol': 1e-4, + 'log': False + } + + def auto_track(self): + results = self.read() + edges = self.compute_dist(results) + results = self.associate(results, edges) + results, occupancy = self.reset_id(results) + results, occupancy = self.smooth(results, occupancy) + self.write(results, occupancy) + + def read(self): + return [] + + def write(self, results, occupancy): + return 0 + + def compute_dist(self, results): + nFrames = len(results) + WINDOW_SIZE = self.WINDOW_SIZE + edges = {} + for start in tqdm(range(0, nFrames - 1), desc='affinity'): + window_size = min(WINDOW_SIZE, nFrames - start) + results_window = results[start:start+window_size] + dimGroups, frames = getDimGroups(results_window) + dist = self._compute_dist(dimGroups, results_window) + res = matchSVT(dist, dimGroups, control=self.svt_args) + xx, yy = np.where(res) + for x, y in zip(xx, yy): + if x >= y:continue + nf0, nf1 = frames[x], frames[y] + ni0, ni1 = x - dimGroups[nf0], y - dimGroups[nf1] + edge = ((nf0+start, ni0), (nf1+start, ni1)) + if edge not in edges: + edges[edge] = [] + edges[edge].append(res[x, y]) + return edges + + def associate(self, results, edges): + WINDOW_SIZE = self.WINDOW_SIZE + connects = list(edges.keys()) + connects.sort(key=lambda x:-sum(edges[x])) + maxid = 0 + frames_of_id = {} + log = print + log = lambda x:x + for (nf0, ni0), (nf1, ni1) in connects: + if abs(nf1 - nf0) > WINDOW_SIZE//2: + continue + # create + id0 = results[nf0][ni0]['id'] + id1 = results[nf1][ni1]['id'] + if id0 == -1 and id1 == -1: + results[nf0][ni0]['id'] = maxid + log('Create person {}'.format(maxid)) + frames_of_id[maxid] = {nf0:ni0, nf1:ni1} + maxid += 1 + # directly assign + if id0 != -1 and id1 == -1: + if nf1 in frames_of_id[id0].keys(): + log('Merge conflict') + results[nf1][ni1]['id'] = id0 + # log('Merge person {}'.format(maxid)) + frames_of_id[id0][nf1] = ni1 + continue + if id0 == -1 and id1 != -1: + results[nf0][ni0]['id'] = id1 + frames_of_id[id1][nf0] = ni0 + continue + if id0 == id1: + continue + # merge + if id0 != id1: + common = frames_of_id[id0].keys() & frames_of_id[id1].keys() + for key in common: # conflict + if frames_of_id[id0][key] == frames_of_id[id1][key]: + pass + else: + break + else: # merge + log('Merge {} to {}'.format(id1, id0)) + for key in frames_of_id[id1].keys(): + results[key][frames_of_id[id1][key]]['id'] = id0 + frames_of_id[id0][key] = frames_of_id[id1][key] + frames_of_id.pop(id1) + continue + log('Conflict; not merged') + return results + + def reset_id(self, results): + mapid = {} + maxid = 0 + occupancy = [] + nFrames = len(results) + for nf, res in enumerate(results): + for info in res: + if info['id'] == -1: + continue + if info['id'] not in mapid.keys(): + mapid[info['id']] = maxid + maxid += 1 + occupancy.append([0 for _ in range(nFrames)]) + pid = mapid[info['id']] + info['id'] = pid + occupancy[pid][nf] = 1 + occupancy = np.array(occupancy) + results, occupancy = self.remove_outlier(results, occupancy) + results, occupancy = self.interpolate(results, occupancy) + return results, occupancy + + def remove_outlier(self, results, occupancy): + nFrames = len(results) + pids = [] + for pid in range(occupancy.shape[0]): + if occupancy[pid].sum() > self.MIN_FRAMES: + pids.append(pid) + occupancy = occupancy[pids] + for nf in range(nFrames): + result = results[nf] + result_filter = [] + for info in result: + if info['id'] == -1 or info['id'] not in pids: + continue + info['id'] = pids.index(info['id']) + result_filter.append(info) + results[nf] = result_filter + return results, occupancy + + def interpolate(self, results, occupancy): + # find missing frames + WINDOW_SIZE = self.WINDOW_SIZE + for pid in range(occupancy.shape[0]): + for nf in range(1, occupancy.shape[1]-1): + if occupancy[pid, nf-1] < 1 or occupancy[pid, nf] > 0: + continue + left = nf - 1 + right = np.where(occupancy[pid, nf+1:])[0] + if len(right) > 0: + right = right.min() + nf + 1 + else: + continue + # find valid (left, right) + # interpolate 3d pose + info_left = [res for res in results[left] if res['id'] == pid][0] + info_right = [res for res in results[right] if res['id'] == pid][0] + for nf_i in range(left+1, right): + weight = 1 - (nf_i - left)/(right - left) + res = self._interpolate(info_left, info_right, weight) + res['id'] = pid + results[nf_i].append(res) + occupancy[pid, nf_i] = pid + return results, occupancy + + def smooth(self, results, occupancy): + return results, occupancy + + def _interpolate(self, info_left, info_right, weight): + return info_left.copy() + +class Track3D(BaseTrack): + def __init__(self, with2d=False, mode='body25', **cfg) -> None: + super().__init__(**cfg) + self.with2d = with2d + self.mode = mode + + def read(self): + k3dpath = join(self.path, 'keypoints3d') + check_path(k3dpath) + filenames = sorted(glob(join(k3dpath, '*.json'))) + if self.with2d: + k2dpath = join(self.path, 'keypoints2d') + check_path(k2dpath) + subs = sorted(os.listdir(k2dpath)) + else: + k2dpath = '' + subs = [] + results = [] + for nf, filename in enumerate(filenames): + basename = os.path.basename(filename) + infos = read_keypoints3d(filename) + for n, info in enumerate(infos): + info['id'] = -1 + info['index'] = n + + results.append(infos) + if self.with2d: + # load 2d keypoints + for nv, sub in enumerate(subs): + k2dname = join(k2dpath, sub, basename) + annots = read_keypoints2d(k2dname, self.mode) + for annot in annots: + pid = annot['id'] + bbox = annot['bbox'] + keypoints = annot['keypoints'] + import ipdb; ipdb.set_trace() + return results + + def write(self, results, mapid): + os.makedirs(self.out, exist_ok=True) + for nf, res in enumerate(results): + outname = join(self.out, 'keypoints3d', '{:06d}.json'.format(nf)) + result = results[nf] + write_keypoints3d(outname, result) + + def _compute_dist(self, dimGroups, results_window): + max_dist = 0.15 + max_dist_step = 0.01 + window_size = len(results_window) + dist = np.eye(dimGroups[-1]) + for i in range(window_size-1): + if len(results_window[i]) == 0: + continue + k3d_i = np.stack([info['keypoints3d'] for info in results_window[i]]) + for j in range(i+1, window_size): + if len(results_window[j]) == 0: + continue + k3d_j = np.stack([info['keypoints3d'] for info in results_window[j]]) + conf = np.sqrt(k3d_i[:, None, :, 3] * k3d_j[None, :, :, 3]) + d_ij = np.linalg.norm(k3d_i[:, None, :, :3] - k3d_j[None, :, :, :3], axis=3) + a_ij = 1 - d_ij / (max_dist + (j-i)*max_dist_step ) + a_ij[a_ij < 0] = 0 + weight =(conf*a_ij).sum(axis=2)/(1e-4 + conf.sum(axis=2)) + dist[dimGroups[i]:dimGroups[i+1], dimGroups[j]:dimGroups[j+1]] = weight + dist[dimGroups[j]:dimGroups[j+1], dimGroups[i]:dimGroups[i+1]] = weight.T + return dist + + def _interpolate(self, info_left, info_right, weight): + kpts_new = info_left['keypoints3d'] * weight + info_right['keypoints3d'] * (1-weight) + res = {'keypoints3d': kpts_new} + return res + +class Track2D(BaseTrack): + def __init__(self, **cfg) -> None: + super().__init__(**cfg) + + def read(self): + filenames = sorted(glob(join(self.path, '*.json'))) + results = [] + for filename in tqdm(filenames, desc='loading'): + result = read_json(filename)['annots'] + for n, info in enumerate(result): + info['id'] = -1 + results.append(result) + return results + + def write(self, results, occupancy): + os.makedirs(self.out, exist_ok=True) + filenames = sorted(glob(join(self.path, '*.json'))) + for nf, res in enumerate(tqdm(results, desc='writing')): + outname = join(self.out, '{:06d}.json'.format(nf)) + result = results[nf] + annots = read_json(filenames[nf]) + annots['annots'] = result + for res in result: + res['personID'] = res.pop('id') + save_annot(outname, annots) + annot = os.path.basename(os.path.dirname(self.out)) + occpath = self.out.replace(annot, 'track') + '.json' + save_json(occpath, occupancy.tolist()) + + def _compute_dist(self, dimGroups, results_window): + window_size = len(results_window) + dist = np.eye(dimGroups[-1]) + for i in range(window_size-1): + if len(results_window[i]) == 0: + continue + bbox_pre = np.stack([info['bbox'] for info in results_window[i]]) + bbox_pre = bbox_pre[:, None] + for j in range(i+1, window_size): + if len(results_window[j]) == 0: + continue + bbox_now = np.stack([info['bbox'] for info in results_window[j]]) + bbox_now = bbox_now[None, :] + areas_pre = (bbox_pre[..., 2] - bbox_pre[..., 0]) * (bbox_pre[..., 3] - bbox_pre[..., 1]) + areas_now = (bbox_now[..., 2] - bbox_now[..., 0]) * (bbox_now[..., 3] - bbox_now[..., 1]) + # 左边界的大值 + xx1 = np.maximum(bbox_pre[..., 0], bbox_now[..., 0]) + yy1 = np.maximum(bbox_pre[..., 1], bbox_now[..., 1]) + # 右边界的小值 + xx2 = np.minimum(bbox_pre[..., 2], bbox_now[..., 2]) + yy2 = np.minimum(bbox_pre[..., 3], bbox_now[..., 3]) + + w = np.maximum(0.0, xx2 - xx1) + h = np.maximum(0.0, yy2 - yy1) + inter = w * h + over = inter / (areas_pre + areas_now - inter) + weight = over + dist[dimGroups[i]:dimGroups[i+1], dimGroups[j]:dimGroups[j+1]] = weight + dist[dimGroups[j]:dimGroups[j+1], dimGroups[i]:dimGroups[i+1]] = weight.T + return dist + + def reset_id(self, results): + results[0].sort(key=lambda x:-(x['bbox'][2]-x['bbox'][0])*(x['bbox'][3]-x['bbox'][1])) + return super().reset_id(results) + + def _interpolate(self, info_left, info_right, weight): + bbox = [info_left['bbox'][i]*weight+info_left['bbox'][i]*(1-weight) for i in range(5)] + kpts_l = info_left['keypoints'] + kpts_r = info_right['keypoints'] + kpts = [] + for nj in range(len(kpts_l)): + if kpts_l[nj][2] < 0.1 or kpts_r[nj][2] < 0.1: + kpts.append([0., 0., 0.]) + else: + kpts.append([kpts_l[nj][i]*weight + kpts_r[nj][i]*(1-weight) for i in range(3)]) + res = {'bbox': bbox, 'keypoints': kpts} + return res + + def smooth(self, results, occupancy): + for pid in range(occupancy.shape[0]): + # the occupancy must be continuous + pass + return results, occupancy \ No newline at end of file diff --git a/easymocap/mytools/reader.py b/easymocap/mytools/reader.py index 36d205b..158a274 100644 --- a/easymocap/mytools/reader.py +++ b/easymocap/mytools/reader.py @@ -2,7 +2,7 @@ @ Date: 2021-04-21 15:19:21 @ Author: Qing Shuai @ LastEditors: Qing Shuai - @ LastEditTime: 2021-06-15 11:30:00 + @ LastEditTime: 2021-06-26 17:37:07 @ FilePath: /EasyMocap/easymocap/mytools/reader.py ''' # function to read data @@ -32,6 +32,8 @@ def read_keypoints3d(filename): # 对于有手的情况，把手的根节点赋值成body25上的点 pose3d[25, :] = pose3d[7, :] pose3d[46, :] = pose3d[4, :] + if pose3d.shape[1] == 3: + pose3d = np.hstack([pose3d, np.ones((pose3d.shape[0], 1))]) res_.append({ 'id': pid, 'keypoints3d': pose3d diff --git a/easymocap/mytools/utils.py b/easymocap/mytools/utils.py index b4ea0bb..7017a83 100644 --- a/easymocap/mytools/utils.py +++ b/easymocap/mytools/utils.py @@ -2,8 +2,8 @@ @ Date: 2021-01-15 11:12:00 @ Author: Qing Shuai @ LastEditors: Qing Shuai - @ LastEditTime: 2021-06-16 14:05:39 - @ FilePath: /EasyMocap/easymocap/mytools/utils.py + @ LastEditTime: 2021-06-25 21:07:29 + @ FilePath: /EasyMocapRelease/easymocap/mytools/utils.py ''' import time import tabulate @@ -42,7 +42,7 @@ class Timer: Timer.records[self.name].append((end-self.start)*1000) if not self.silent: t = (end - self.start)*1000 - if t > 10000: + if t > 1000: print('-> [{:20s}]: {:5.1f}s'.format(self.name, t/1000)) elif t > 1e3*60*60: print('-> [{:20s}]: {:5.1f}min'.format(self.name, t/1e3/60)) diff --git a/easymocap/pipeline/weight.py b/easymocap/pipeline/weight.py index 8810a8c..2ebfba4 100644 --- a/easymocap/pipeline/weight.py +++ b/easymocap/pipeline/weight.py @@ -2,11 +2,16 @@ @ Date: 2021-04-13 20:12:58 @ Author: Qing Shuai @ LastEditors: Qing Shuai - @ LastEditTime: 2021-05-27 17:04:47 - @ FilePath: /EasyMocap/easymocap/pipeline/weight.py + @ LastEditTime: 2021-06-25 22:14:58 + @ FilePath: /EasyMocapRelease/easymocap/pipeline/weight.py ''' -def load_weight_shape(opts): - weight = {'s3d': 1., 'reg_shapes': 5e-3} +def load_weight_shape(model, opts): + if model in ['smpl', 'smplh', 'smplx']: + weight = {'s3d': 1., 'reg_shapes': 5e-3} + elif model == 'mano': + weight = {'s3d': 1e2, 'reg_shapes': 5e-5} + else: + raise NotImplementedError for key in opts.keys(): if key in weight.keys(): weight[key] = opts[key] @@ -15,8 +20,8 @@ def load_weight_shape(opts): def load_weight_pose(model, opts): if model == 'smpl': weight = { - 'k3d': 1., 'reg_poses_zero': 1e-2, 'smooth_body': 5e-1, - 'smooth_poses': 1e-1, 'reg_poses': 1e-3, + 'k3d': 1., 'reg_poses_zero': 1e-2, 'smooth_body': 5e0, + 'smooth_poses': 1e0, 'reg_poses': 1e-3, 'k2d': 1e-4 } elif model == 'smplh': @@ -37,9 +42,14 @@ def load_weight_pose(model, opts): } elif model == 'mano': weight = { - 'k3d': 1e2, 'k2d': 1e-3, - 'reg_poses': 1e-3, 'smooth_body': 1e2 + 'k3d': 1e2, 'k2d': 2e-3, + 'reg_poses': 1e-3, 'smooth_body': 1e2, + # 'collision': 1 # If the frame number is too large (more than 1000), then GPU oom } + # weight = { + # 'k3d': 1., 'k2d': 1e-4, + # 'reg_poses': 1e-4, 'smooth_body': 0 + # } else: print(model) raise NotImplementedError diff --git a/scripts/postprocess/eval_shelf.py b/scripts/postprocess/eval_shelf.py new file mode 100644 index 0000000..17e7e9b --- /dev/null +++ b/scripts/postprocess/eval_shelf.py @@ -0,0 +1,289 @@ +''' + @ Date: 2020-12-01 22:14:11 + @ Author: Qing Shuai + @ LastEditors: Qing Shuai + @ LastEditTime: 2021-05-30 21:33:40 + @ FilePath: /EasyMocap/scripts/postprocess/eval_shelf.py +''' +import os +import sys +from os.path import join +import re +import json +import time +import scipy.io as scio +import numpy as np +from tqdm import tqdm + +def save_json(output, json_path): + os.system('mkdir -p {}'.format(os.path.dirname(json_path))) + with open(json_path, 'w') as f: + json.dump(output, f, indent=4) + +def is_right(model_start_point, model_end_point, gt_strat_point, gt_end_point, alpha=0.5): + bone_lenth = np.linalg.norm ( gt_end_point - gt_strat_point ) + start_difference = np.linalg.norm ( gt_strat_point - model_start_point ) + end_difference = np.linalg.norm ( gt_end_point - model_end_point ) + return ((start_difference + end_difference) / 2) <= alpha * bone_lenth + +def openpose2shelf3D(pose3d, score): + """ + transform coco order(our method output) 3d pose to shelf dataset order with interpolation + :param pose3d: np.array with shape nJx3 + :return: 3D pose in shelf order with shape 14x3 + """ + shelf_pose = np.zeros ( (14, 3) ) + shelf_score = np.zeros ( (14, 1) ) + + # coco2shelf = np.array ( [16, 14, 12, 11, 13, 15, 10, 8, 6, 5, 7, 9] ) + openpose2shelf = np.array([11, 10, 9, 12, 13, 14, 4, 3, 2, 5, 6, 7]) + shelf_pose[0: 12] += pose3d[openpose2shelf] + shelf_score[0: 12] += score[openpose2shelf] + if True: + shelf_pose[12] = pose3d[1] # Use middle of shoulder to init + shelf_pose[13] = pose3d[0] # use nose to init + shelf_pose[13] = shelf_pose[12] + (shelf_pose[13] - shelf_pose[12]) * np.array ( [0.75, 0.75, 1.5] ) + shelf_pose[12] = shelf_pose[12] + (pose3d[0] - shelf_pose[12]) * np.array ( [1. / 2., 1. / 2., 1. / 2.] ) + shelf_score[12] = score[0]*score[1] + shelf_score[13] = score[0]*score[1] + else: + shelf_pose[12] = pose3d[1] + shelf_pose[13] = pose3d[0] + return shelf_pose, shelf_score + +def convert_openpose_shelf(keypoints3d): + shelf15 = np.zeros((15, 4)) + openpose2shelf = np.array([11, 10, 9, 12, 13, 14, 4, 3, 2, 5, 6, 7, 1, 0, 8]) + shelf15 = keypoints3d[openpose2shelf].copy() + # interp head + faceDir = np.cross(shelf15[12, :3] - shelf15[14, :3], shelf15[8, :3] - shelf15[9, :3]) + faceDir = faceDir/np.linalg.norm(faceDir) + zDir = np.array([0., 0., 1.]) + shoulderCenter = (keypoints3d[2, :3] + keypoints3d[5, :3])/2. + # headCenter = (keypoints3d[15, :3] + keypoints3d[16, :3])/2. + headCenter = (keypoints3d[17, :3] + keypoints3d[18, :3])/2. + + shelf15[12, :3] = shoulderCenter + (headCenter - shoulderCenter) * 0.5 + shelf15[13, :3] = shelf15[12, :3] + faceDir * 0.125 + zDir * 0.145 + return shelf15 + +def convert_openpose_shelf1(keypoints3d): + shelf15 = np.zeros((15, 4)) + openpose2shelf = np.array([11, 10, 9, 12, 13, 14, 4, 3, 2, 5, 6, 7, 1, 0, 8]) + shelf15 = keypoints3d[openpose2shelf].copy() + # interp head + faceDir = np.cross(keypoints3d[1, :3] - keypoints3d[8, :3], keypoints3d[2, :3] - shelf15[5, :3]) + faceDir = faceDir/np.linalg.norm(faceDir) + + upDir = keypoints3d[1, :3] - keypoints3d[8, :3] + upDir = upDir/np.linalg.norm(upDir) + + shoulderCenter = keypoints3d[1, :3] + ear = (keypoints3d[17, :3] + keypoints3d[18, :3])/2 - keypoints3d[1, :3] + eye = (keypoints3d[15, :3] + keypoints3d[16, :3])/2 - keypoints3d[1, :3] + nose = keypoints3d[0, :3] - keypoints3d[1, :3] + head = (ear + eye + nose)/3. + noseLen = np.linalg.norm(head) + noseDir = head / noseLen + headDir = (noseDir * 2 + upDir) + headDir = headDir / np.linalg.norm(headDir) + + neck = shoulderCenter + noseLen*headDir * 0.5 + + shelf15[12, :3] = neck + shelf15[13, :3] = neck + headDir * noseLen * 0.8 + return shelf15 + +def convert_shelf_shelfgt(keypoints): + gt_hip = (keypoints[2] + keypoints[3]) / 2 + gt = np.vstack((keypoints, gt_hip)) + return gt + +def vectorize_distance(a, b): + """ + Calculate euclid distance on each row of a and b + :param a: Nx... np.array + :param b: Mx... np.array + :return: MxN np.array representing correspond distance + """ + N = a.shape[0] + a = a.reshape ( N, -1 ) + M = b.shape[0] + b = b.reshape ( M, -1 ) + a2 = np.tile ( np.sum ( a ** 2, axis=1 ).reshape ( -1, 1 ), (1, M) ) + b2 = np.tile ( np.sum ( b ** 2, axis=1 ), (N, 1) ) + dist = a2 + b2 - 2 * (a @ b.T) + return np.sqrt ( dist ) + +def distance(a, b, score): + # a: (N, J, 3) + # b: (M, J, 3) + # score: (M, J, 1) + # return: (M, N) + a = a[None, :, :, :] + b = b[:, None, :, :] + score = score[:, None, :, 0] + diff = np.sum((a - b)**2, axis=3)*score + dist = diff.sum(axis=2)/score.sum(axis=2) + return np.sqrt(dist) + +def _readResult(filename, isA4d): + import json + with open(filename, "r") as file: + datas = json.load(file) + res_ = [] + for data in datas: + trackId = data['id'] + keypoints3d = np.array(data['keypoints3d']) + if (keypoints3d[:, 3]>0).sum() > 1: + res_.append({'id':trackId, 'keypoints3d': keypoints3d}) + if isA4d: + # association4d 的关节顺序和正常的定义不一样 + for r in res_: + r['keypoints3d'] = r['keypoints3d'][[4, 1, 5, 9, 13, 6, 10, 14, 0, 2, 7, 11, 3, 8, 12], :] + return res_ + +def readResult(filePath, range_=None, isA4d=None): + res = {} + if range_ is None: + from glob import glob + filelists = glob(join(filePath, '*.txt')) + range_ = [i for i in range(len(filelists))] + if isA4d is None: + isA4d = args.a4d + for imgId in tqdm(range_): + res[imgId] = _readResult(join(filePath, '{:06d}.json'.format(imgId)), isA4d) + return res + +class ShelfGT: + def __init__(self, actor3D) -> None: + self.actor3D = actor3D + self.actor3D = self.actor3D[:3] + + def __getitem__(self, index): + results = [] + for pid in range(len(self.actor3D)): + gt_pose = self.actor3D[pid][index-2][0] + if gt_pose.shape == (1, 0) or gt_pose.shape == (0, 0): + continue + keypoints3d = convert_shelf_shelfgt(gt_pose) + results.append({'id': pid, 'keypoints3d': keypoints3d}) + return results + +def write_to_csv(filename, results, id_wise=True): + keys = [key for key in results[0].keys() if isinstance(results[0][key], float)] + if id_wise: + ids = list(set([res['id'] for res in results])) + header = [''] + ['{:s}'.format(key.replace(' ', '')) for key in keys] + contents = [] + if id_wise: + for pid in ids: + content = ['{}'.format(pid)] + for key in keys: + vals = [res[key] for res in results if res['id'] == pid] + content.append('{:.3f}'.format(sum(vals)/len(vals))) + contents.append(content) + # 计算平均值 + content = ['Mean'] + for i, key in enumerate(keys): + content.append('{:.3f}'.format(sum([float(con[i+1]) for con in contents])/len(ids))) + contents.append(content) + else: + content = ['Mean'] + for key in keys: + content.append('{:.3f}'.format(sum([res[key] for res in results])/len(results))) + contents.append(content) + import tabulate + print(tabulate.tabulate(contents, header, tablefmt='fancy_grid')) + print(tabulate.tabulate(contents, header, tablefmt='fancy_grid'), file=open(filename.replace('.csv', '.txt'), 'w')) + + with open(filename, 'w') as f: + # 写入头 + header = list(results[0].keys()) + f.write(','.join(header) + '\n') + for res in results: + f.write(','.join(['{}'.format(res[key]) for key in header]) + '\n') + +def evaluate(actor3D, range_, out): + shelfgt = ShelfGT(actor3D) + check_result = np.zeros ( (len ( actor3D[0] ), len ( actor3D ), 10), dtype=np.int32 ) + result = readResult(out, range_) + bones = [[0, 1], [1, 2], [3, 4], [4, 5], [6, 7], [7, 8], [9, 10], [10, 11], [12, 13], [12, 14]] + start = [ 9, 8, 10, 7, 3, 2, 4, 1, 12, 12,] + end = [10, 7, 11, 6, 4, 1, 5, 0, 13, 14] + names = ["Left Upper Arm", "Right Upper Arm", "Left Lower Arm", "Right Lower Arm", "Left Upper Leg", "Right Upper Leg", "Left Lower Leg", "Right Lower Leg", "Head", "Torso" ] + + results = [] + for img_id in tqdm(range_): + # 转化成model_poses + ests = [] + for res in result[img_id]: + ests.append({'id': res['id'], 'keypoints3d': convert_openpose_shelf1(res['keypoints3d'])}) + gts = shelfgt[img_id] + if len(gts) < 1: + continue + # 匹配最近的 + kpts_gt = np.stack([v['keypoints3d'] for v in gts]) + kpts_dt = np.stack([v['keypoints3d'] for v in ests]) + distances = np.linalg.norm(kpts_gt[:, None, :, :3] - kpts_dt[None, :, :, :3], axis=-1) + conf = (kpts_gt[:, None, :, -1] > 0) * (kpts_dt[None, :, :, -1] > 0) + dist = (distances * conf).sum(axis=-1)/conf.sum(axis=-1) + # 贪婪的匹配 + ests_new = [] + for igt, gt in enumerate(gts): + bestid = np.argmin(dist[igt]) + ests_new.append(ests[bestid]) + ests = ests_new + # 计算误差 + for i, data in enumerate(gts): + kpts_gt = data['keypoints3d'] + kpts_est = ests[i]['keypoints3d'] + # 计算各种误差，存成字典 + da = np.linalg.norm(kpts_gt[start, :3] - kpts_est[start, :3], axis=1) + db = np.linalg.norm(kpts_gt[end, :3] - kpts_est[end, :3], axis=1) + l = np.linalg.norm(kpts_gt[start, :3] - kpts_gt[end, :3], axis=1) + isright = 1.0*((da + db) < l) + if args.joint: + res = {name: isright[i] for i, name in enumerate(names)} + else: + res = {} + res['Mean'] = isright.mean() + res['nf'] = img_id + res['id'] = data['id'] + results.append(res) + write_to_csv(join(out, '..', 'report.csv'), results) + return 0 + +if __name__ == '__main__': + + import argparse + + parser = argparse.ArgumentParser () + parser.add_argument('--out', type=str, default='output/') + parser.add_argument('--gt_path', type=str, default='config/evaluation/actorsGT_shelf.mat') + parser.add_argument('--setting', type=str, default='shelf') + parser.add_argument('--a4d', action='store_true') + parser.add_argument('--joint', action='store_true') + + args = parser.parse_args () + if args.setting == 'shelf': + test_range = range ( 302, 602) + # test_range = range (2000, 3200) + elif args.setting == 'campus': + test_range = [i for i in range ( 350, 471 )] + [i for i in range ( 650, 751 )] + else: + raise NotImplementedError + + actorsGT = scio.loadmat (args.gt_path) + test_actor3D = actorsGT['actor3D'][0] + if False: + valid = np.zeros((3200, 4)) + for nf in range(3200): + for pid in range(4): + if test_actor3D[pid][nf].item().shape[0] == 14: + valid[nf, pid] = 1 + import matplotlib.pyplot as plt + plt.plot(valid.sum(axis=1)) + plt.show() + import ipdb; ipdb.set_trace() + evaluate(test_actor3D, test_range, args.out)