from operator import imod import numpy as np from tqdm import tqdm from os.path import join from easymocap.dataset.mv1pmf_mirror import ImageFolderMirror as ImageFolder from easymocap.mytools import Timer from easymocap.smplmodel import load_model, merge_params, select_nf from easymocap.estimator import SPIN, init_with_spin from easymocap.pipeline.mirror import multi_stage_optimize def demo_1v1p1f_smpl_mirror(path, body_model, spin_model, args): "Optimization for single image" # 0. construct the dataset dataset = ImageFolder(path, out=args.out, kpts_type=args.body) if args.gtK: dataset.gtK = True dataset.load_gt_cameras() start, end = args.start, min(args.end, len(dataset)) for nf in tqdm(range(start, end, args.step), desc='Optimizing'): image, annots = dataset[nf] if len(annots) < 2: continue annots = annots[:2] camera = dataset.camera(nf) # initialize the SMPL parameters body_params_all = [] bboxes, keypoints2d, pids = [], [], [] for i, annot in enumerate(annots): assert annot['id'] == i, (i, annot['id']) result = init_with_spin(body_model, spin_model, image, annot['bbox'], annot['keypoints'], camera) body_params_all.append(result['body_params']) bboxes.append(annot['bbox']) keypoints2d.append(annot['keypoints']) pids.append(annot['id']) bboxes = np.vstack(bboxes) keypoints2d = np.stack(keypoints2d) body_params = merge_params(body_params_all) # bboxes: (nViews(2), 1, 5); keypoints2d: (nViews(2), 1, nJoints, 3) bboxes = bboxes[:, None] keypoints2d = keypoints2d[:, None] if args.normal: normal = dataset.normal(nf)[None, :, :] else: normal = None body_params = multi_stage_optimize(body_model, body_params, bboxes, keypoints2d, Pall=camera['P'], normal=normal, args=args) vertices = body_model(return_verts=True, return_tensor=False, **body_params) keypoints = body_model(return_verts=False, return_tensor=False, **body_params) write_data = [{'id': pids[i], 'keypoints3d': keypoints[i]} for i in range(len(pids))] # write out the results dataset.write_keypoints3d(write_data, nf) for i in range(len(pids)): write_data[i].update(select_nf(body_params, i)) if args.vis_smpl: # render the results render_data = {pids[i]: { 'vertices': vertices[i], 'faces': body_model.faces, 'vid': 0, 'name': 'human_{}'.format(pids[i])} for i in range(len(pids))} dataset.vis_smpl(render_data, image, camera, nf) dataset.write_smpl(write_data, nf) def demo_1v1pmf_smpl_mirror(path, body_model, spin_model, args): subs = args.sub assert len(subs) > 0 # 遍历所有文件夹 for sub in subs: dataset = ImageFolder(path, subs=[sub], out=args.out, kpts_type=args.body) start, end = args.start, min(args.end, len(dataset)) frames = list(range(start, end, args.step)) nFrames = len(frames) pids = [0, 1] body_params_all = {pid:[None for nf in frames] for pid in pids} bboxes = {pid:[None for nf in frames] for pid in pids} keypoints2d = {pid:[None for nf in frames] for pid in pids} for nf in tqdm(frames, desc='loading'): image, annots = dataset[nf] # 这个时候如果annots不够 不能够跳过了，需要进行补全 camera = dataset.camera(nf) # 初始化每个人的SMPL参数 for i, annot in enumerate(annots): pid = annot['id'] if pid not in pids: continue result = init_with_spin(body_model, spin_model, image, annot['bbox'], annot['keypoints'], camera) body_params_all[pid][nf-start] = result['body_params'] bboxes[pid][nf-start] = annot['bbox'] keypoints2d[pid][nf-start] = annot['keypoints'] # stack [p1f1, p1f2, p1f3, ..., p1fn, p2f1, p2f2, p2f3, ..., p2fn] # TODO:for missing bbox body_params = merge_params([merge_params(body_params_all[pid]) for pid in pids]) # bboxes: (nViews, nFrames, 5) bboxes = np.stack([np.stack(bboxes[pid]) for pid in pids]) # keypoints: (nViews, nFrames, nJoints, 3) keypoints2d = np.stack([np.stack(keypoints2d[pid]) for pid in pids]) # optimize P = dataset.camera(start)['P'] if args.normal: normal = dataset.normal_all(start=start, end=end) else: normal = None body_params = multi_stage_optimize(body_model, body_params, bboxes, keypoints2d, Pall=P, normal=normal, args=args) # write vertices = body_model(return_verts=True, return_tensor=False, **body_params) keypoints = body_model(return_verts=False, return_tensor=False, **body_params) dataset.no_img = not args.vis_smpl for nf in tqdm(frames, desc='rendering'): idx = nf - start write_data = [{'id': pids[i], 'keypoints3d': keypoints[i*nFrames+idx]} for i in range(len(pids))] dataset.write_keypoints3d(write_data, nf) for i in range(len(pids)): write_data[i].update(select_nf(body_params, i*nFrames+idx)) dataset.write_smpl(write_data, nf) # 保存结果 if args.vis_smpl: image, annots = dataset[nf] camera = dataset.camera(nf) render_data = {pids[i]: { 'vertices': vertices[i*nFrames+idx], 'faces': body_model.faces, 'vid': 0, 'name': 'human_{}'.format(pids[i])} for i in range(len(pids))} dataset.vis_smpl(render_data, image, camera, nf) if __name__ == "__main__": from easymocap.mytools import load_parser, parse_parser parser = load_parser() parser.add_argument('--skel', type=str, default=None, help='path to keypoints3d') parser.add_argument('--direct', action='store_true') parser.add_argument('--video', action='store_true') parser.add_argument('--gtK', action='store_true') parser.add_argument('--normal', action='store_true', help='set to use the normal of the mirror') args = parse_parser(parser) helps = ''' Demo code for single view and one person with mirror: - Input : {}: [{}] - Output: {} - Body : {} => {}, {} '''.format(args.path, ', '.join(args.sub), args.out, args.model, args.gender, args.body) print(helps) with Timer('Loading {}, {}'.format(args.model, args.gender)): body_model = load_model(args.gender, model_type=args.model) with Timer('Loading SPIN'): spin_model = SPIN( SMPL_MEAN_PARAMS='data/models/smpl_mean_params.npz', checkpoint='data/models/spin_checkpoint.pt', device=body_model.device) if args.video: demo_1v1pmf_smpl_mirror(args.path, body_model, spin_model, args) else: demo_1v1p1f_smpl_mirror(args.path, body_model, spin_model, args)