def create_model(args, dataset, poses_valid_2d): filter_widths = [int(x) for x in args.architecture.split(",")] adj = adj_mx_from_skeleton(dataset.skeleton()) if not args.disable_optimizations and args.stride == 1: # Use optimized model for single-frame predictions model_pos_train = SpatioTemporalModelOptimized1f( adj, poses_valid_2d[0].shape[-2], poses_valid_2d[0].shape[-1], dataset.skeleton().num_joints(), filter_widths=filter_widths, causal=args.causal, dropout=args.dropout, channels=args.channels) else: # When incompatible settings are detected (stride > 1, dense filters, or disabled optimization) fall back to normal model model_pos_train = SpatioTemporalModel(adj, poses_valid_2d[0].shape[-2], poses_valid_2d[0].shape[-1], dataset.skeleton().num_joints(), filter_widths=filter_widths, causal=args.causal, dropout=args.dropout, channels=args.channels) model_pos = SpatioTemporalModel(adj, poses_valid_2d[0].shape[-2], poses_valid_2d[0].shape[-1], dataset.skeleton().num_joints(), filter_widths=filter_widths, causal=args.causal, dropout=args.dropout, channels=args.channels) receptive_field = model_pos.receptive_field() print("INFO: Receptive field: {} frames".format(receptive_field)) pad = (receptive_field - 1) // 2 # padding on each side if args.causal: print("INFO: Using causal convolutions") causal_shift = pad else: causal_shift = 0 model_params = 0 for parameter in model_pos.parameters(): model_params += parameter.numel() print("INFO: Trainable parameter count: ", model_params) return model_pos_train, model_pos, pad, causal_shift
def main(args): print('==> Using settings {}'.format(args)) # print('==> Loading dataset...') # dataset_path = path.join('data', 'data_3d_' + args.dataset + '.npz') # if args.dataset == 'h36m': # from common.h36m_dataset import Human36mDataset, TRAIN_SUBJECTS, TEST_SUBJECTS # dataset = Human36mDataset(dataset_path) # subjects_train = TRAIN_SUBJECTS # subjects_test = TEST_SUBJECTS # else: # raise KeyError('Invalid dataset') # print('==> Preparing data...') # dataset = read_3d_data(dataset) # # print('==> Loading 2D detections...') # keypoints = create_2d_data(path.join('data', 'data_2d_' + args.dataset + '_' + args.keypoints + '.npz'), dataset) # # action_filter = None if args.actions == '*' else args.actions.split(',') # if action_filter is not None: # action_filter = map(lambda x: dataset.define_actions(x)[0], action_filter) # print('==> Selected actions: {}'.format(action_filter)) # stride = args.downsample cudnn.benchmark = True device = torch.device("cuda") # Create model print("==> Creating model...") p_dropout = (None if args.dropout == 0.0 else args.dropout) # adj = adj_mx_from_skeleton(dataset.skeleton()) adj = adj_mx_from_skeleton() print(adj) # model_pos = SemGCN(adj, args.hid_dim, num_layers=args.num_layers, p_dropout=p_dropout, # nodes_group=dataset.skeleton().joints_group() if args.non_local else None).to(device) model_pos = SemGCN(adj, args.hid_dim, num_layers=args.num_layers, p_dropout=p_dropout, nodes_group=None).to(device) print("==> Total parameters: {:.2f}M".format( sum(p.numel() for p in model_pos.parameters()) / 1000000.0)) criterion = nn.MSELoss(reduction='mean').to(device) optimizer = torch.optim.Adam(model_pos.parameters(), lr=args.lr) # Optionally resume from a checkpoint if args.resume or args.evaluate: ckpt_path = (args.resume if args.resume else args.evaluate) if path.isfile(ckpt_path): print("==> Loading checkpoint '{}'".format(ckpt_path)) ckpt = torch.load(ckpt_path) start_epoch = ckpt['epoch'] error_best = ckpt['error'] glob_step = ckpt['step'] lr_now = ckpt['lr'] model_pos.load_state_dict(ckpt['state_dict']) optimizer.load_state_dict(ckpt['optimizer']) print("==> Loaded checkpoint (Epoch: {} | Error: {})".format( start_epoch, error_best)) if args.resume: ckpt_dir_path = path.dirname(ckpt_path) logger = Logger(path.join(ckpt_dir_path, 'log.txt'), resume=True) else: raise RuntimeError( "==> No checkpoint found at '{}'".format(ckpt_path)) else: start_epoch = 0 error_best = None glob_step = 0 lr_now = args.lr ckpt_dir_path = args.checkpoint + '/' + 'test1' if not path.exists(ckpt_dir_path): os.makedirs(ckpt_dir_path) print('==> Making checkpoint dir: {}'.format(ckpt_dir_path)) logger = Logger(os.path.join(ckpt_dir_path, 'log.txt')) logger.set_names( ['epoch', 'lr', 'loss_train', 'error_eval_p1', 'error_eval_p2']) if args.evaluate: print('==> Evaluating...') # if action_filter is None: # action_filter = dataset.define_actions() # # errors_p1 = np.zeros(len(action_filter)) # errors_p2 = np.zeros(len(action_filter)) # for i, action in enumerate(action_filter): # poses_valid, poses_valid_2d, actions_valid = fetch(subjects_test, dataset, keypoints, [action], stride) # valid_loader = DataLoader(PoseGenerator(poses_valid, poses_valid_2d, actions_valid), # batch_size=args.batch_size, shuffle=False, # num_workers=args.num_workers, pin_memory=True) # poses_valid, poses_valid_2d, actions_valid = fetch(subjects_test, dataset, keypoints, [action], stride) valid_loader = DataLoader(PoseGenerator(opt='val'), batch_size=args.batch_size, shuffle=False, num_workers=args.num_workers, pin_memory=True) errors_p1, errors_p2 = evaluate(valid_loader, model_pos, device) print('Protocol #1 (MPJPE) action-wise average: {:.2f} (mm)'.format( np.mean(errors_p1))) print('Protocol #2 (P-MPJPE) action-wise average: {:.2f} (mm)'.format( np.mean(errors_p2))) # print('Protocol #1 (MPJPE) action-wise average: {:.2f} (mm)'.format(np.mean(errors_p1).item())) # print('Protocol #2 (P-MPJPE) action-wise average: {:.2f} (mm)'.format(np.mean(errors_p2).item())) exit(0) # poses_train, poses_train_2d, actions_train = fetch(subjects_train, dataset, keypoints, action_filter, stride) # print('2d data shape: ', len(poses_train_2d)) # print('3d data shape: ', len(poses_train)) # print('action data shape: ', len(actions_train)) # train_loader = DataLoader(PoseGenerator(poses_train, poses_train_2d, actions_train), batch_size=args.batch_size, # shuffle=True, num_workers=args.num_workers, pin_memory=True) train_loader = DataLoader(PoseGenerator(opt='train'), batch_size=args.batch_size, shuffle=True, num_workers=args.num_workers, pin_memory=True) # poses_valid, poses_valid_2d, actions_valid = fetch(subjects_test, dataset, keypoints, action_filter, stride) # valid_loader = DataLoader(PoseGenerator(poses_valid, poses_valid_2d, actions_valid), batch_size=args.batch_size, # shuffle=False, num_workers=args.num_workers, pin_memory=True) valid_loader = DataLoader(PoseGenerator(opt='val'), batch_size=args.batch_size, shuffle=False, num_workers=args.num_workers, pin_memory=True) for epoch in range(start_epoch, args.epochs): print('\nEpoch: %d | LR: %.8f' % (epoch + 1, lr_now)) # Train for one epoch epoch_loss, lr_now, glob_step = train(train_loader, model_pos, criterion, optimizer, device, args.lr, lr_now, glob_step, args.lr_decay, args.lr_gamma, max_norm=args.max_norm) # Evaluate error_eval_p1, error_eval_p2 = evaluate(valid_loader, model_pos, device) # Update log file logger.append( [epoch + 1, lr_now, epoch_loss, error_eval_p1, error_eval_p2]) # Save checkpoint if error_best is None or error_best > error_eval_p1: error_best = error_eval_p1 save_ckpt( { 'epoch': epoch + 1, 'lr': lr_now, 'step': glob_step, 'state_dict': model_pos.state_dict(), 'optimizer': optimizer.state_dict(), 'error': error_eval_p1 }, ckpt_dir_path, suffix='best') if (epoch + 1) % args.snapshot == 0: save_ckpt( { 'epoch': epoch + 1, 'lr': lr_now, 'step': glob_step, 'state_dict': model_pos.state_dict(), 'optimizer': optimizer.state_dict(), 'error': error_eval_p1 }, ckpt_dir_path) logger.close() logger.plot(['loss_train', 'error_eval_p1']) savefig(path.join(ckpt_dir_path, 'log.eps')) return
def main(args): human36m_data_path = os.path.join('data', 'data_3d_' + "h36m" + '.npz') MUCO3DHP_path = "/home/lgh/data1/multi3Dpose/muco-3dhp/output/unaugmented_set_001" hid_dim = 128 num_layers = 4 non_local = True lr = 1.0e-3 epochs = 30 _lr_decay = 100000 lr_gamma = 0.96 max_norm = True num_workers = 8 snapshot = 5 batch_size = 64 print('==> Loading multi-person dataset...') #human36m_dataset_path = path.join(human36m_data_path) data_2d, data_3d, img_name, feature_mutual = get_MUCO3DHP_data( MUCO3DHP_path, args) ## N * (M*17) * 2 N * (M*17) * 3 numpy print(img_name[1]) ax = plt.subplot(111, projection='3d') #ax.scatter(data_3d[0, 0::4, 0], data_3d[0, 0::4, 1], data_3d[0, 0::4, 2]) viz.show3Dpose(data_3d[1, :, :], ax) plt.show() person_num = data_2d.shape[1] / 17 dataset = CMUPanoDataset(human36m_data_path, person_num) ### divide into trainsets and testsets 4/5 and 1/5 num = len(data_2d) train_num = num * 4 / 5 cudnn.benchmark = True device = torch.device("cuda") adj, adj_mutual = adj_mx_from_skeleton(dataset.skeleton(), person_num) #ok model_pos = MultiSemGCN(adj, adj_mutual, person_num, hid_dim, num_layers=num_layers, nodes_group=dataset.skeleton().joints_group() if non_local else None).to(device) #ok criterion = nn.MSELoss(reduction='mean').to(device) optimizer = torch.optim.Adam(model_pos.parameters(), lr=lr) start_epoch = 0 error_best = None glob_step = 0 lr_now = lr ckpt_dir_path = os.path.join( 'checkpoint_multi', datetime.datetime.now().isoformat() + "_l_%04d_hid_%04d_e_%04d_non_local_%d" % (num_layers, hid_dim, epochs, non_local)) if not os.path.exists(ckpt_dir_path): os.makedirs(ckpt_dir_path) print('=> Making checkpoint dir: {}'.format(ckpt_dir_path)) logger = Logger(os.path.join(ckpt_dir_path, 'log.txt')) logger.set_names( ['epoch', 'lr', 'loss_train', 'error_eval_p1', 'error_eval_p2']) train_loader = DataLoader(PoseGenerator(data_3d[:train_num], data_2d[:train_num], feature_mutual[:train_num]), batch_size=batch_size, shuffle=True, num_workers=num_workers, pin_memory=True) valid_loader = DataLoader(PoseGenerator(data_3d[train_num:], data_2d[train_num:], feature_mutual[train_num:]), batch_size=batch_size, shuffle=False, num_workers=num_workers, pin_memory=True) writer = SummaryWriter() for epoch in range(start_epoch, epochs): # Train for one epoch print('\nEpoch: %d | LR: %.8f' % (epoch + 1, lr_now)) epoch_loss, lr_now, glob_step = train(train_loader, model_pos, criterion, optimizer, device, lr, lr_now, glob_step, _lr_decay, lr_gamma, max_norm=max_norm) writer.add_scalar('epoch_loss', epoch_loss, epoch) # Evaluate error_eval_p1, error_eval_p2 = evaluate(valid_loader, model_pos, device) # Update log file logger.append( [epoch + 1, lr_now, epoch_loss, error_eval_p1, error_eval_p2]) # Save checkpoint if error_best is None or error_best > error_eval_p1: error_best = error_eval_p1 save_ckpt( { 'epoch': epoch + 1, 'lr': lr_now, 'step': glob_step, 'state_dict': model_pos.state_dict(), 'optimizer': optimizer.state_dict(), 'error': error_eval_p1 }, ckpt_dir_path, suffix='best') if (epoch + 1) % snapshot == 0: save_ckpt( { 'epoch': epoch + 1, 'lr': lr_now, 'step': glob_step, 'state_dict': model_pos.state_dict(), 'optimizer': optimizer.state_dict(), 'error': error_eval_p1 }, ckpt_dir_path) logger.close() writer.close() logger.plot(['loss_train', 'error_eval_p1']) savefig(os.path.join(ckpt_dir_path, 'log.eps'))
def main(args): print('==> Using settings {}'.format(args)) convm = torch.zeros(3, 17, 17, dtype=torch.float) print('==> Loading dataset...') dataset_path = path.join('data', 'data_3d_' + args.dataset + '.npz') if args.dataset == 'h36m': from common.h36m_dataset import Human36mDataset dataset = Human36mDataset(dataset_path) else: raise KeyError('Invalid dataset') print('==> Preparing data...') dataset = read_3d_data(dataset) print('==> Loading 2D detections...') keypoints = create_2d_data( path.join('data', 'data_2d_' + args.dataset + '_' + args.keypoints + '.npz'), dataset) cudnn.benchmark = True device = torch.device("cuda") # Create model print("==> Creating model...") if args.architecture == 'linear': from models.linear_model import LinearModel, init_weights num_joints = dataset.skeleton().num_joints() model_pos = LinearModel(num_joints * 2, (num_joints - 1) * 3).to(device) model_pos.apply(init_weights) elif args.architecture == 'gcn': from models.sem_gcn import SemGCN from common.graph_utils import adj_mx_from_skeleton p_dropout = (None if args.dropout == 0.0 else args.dropout) adj = adj_mx_from_skeleton(dataset.skeleton()) model_pos = SemGCN(convm, adj, args.hid_dim, num_layers=args.num_layers, p_dropout=p_dropout, nodes_group=dataset.skeleton().joints_group() if args.non_local else None).to(device) else: raise KeyError('Invalid model architecture') print("==> Total parameters: {:.2f}M".format( sum(p.numel() for p in model_pos.parameters()) / 1000000.0)) # Resume from a checkpoint ckpt_path = args.evaluate if path.isfile(ckpt_path): print("==> Loading checkpoint '{}'".format(ckpt_path)) ckpt = torch.load(ckpt_path) start_epoch = ckpt['epoch'] error_best = ckpt['error'] model_pos.load_state_dict(ckpt['state_dict']) print("==> Loaded checkpoint (Epoch: {} | Error: {})".format( start_epoch, error_best)) else: raise RuntimeError("==> No checkpoint found at '{}'".format(ckpt_path)) print('==> Rendering...') poses_2d = keypoints[args.viz_subject][args.viz_action] out_poses_2d = poses_2d[args.viz_camera] out_actions = [args.viz_camera] * out_poses_2d.shape[0] poses_3d = dataset[args.viz_subject][args.viz_action]['positions_3d'] assert len(poses_3d) == len(poses_2d), 'Camera count mismatch' out_poses_3d = poses_3d[args.viz_camera] ground_truth = dataset[args.viz_subject][args.viz_action]['positions_3d'][ args.viz_camera].copy() input_keypoints = out_poses_2d.copy() render_loader = DataLoader(PoseGenerator([out_poses_3d], [out_poses_2d], [out_actions]), batch_size=args.batch_size, shuffle=False, num_workers=args.num_workers, pin_memory=True) prediction = evaluate(render_loader, model_pos, device, args.architecture)[0] # Invert camera transformation cam = dataset.cameras()[args.viz_subject][args.viz_camera] prediction = camera_to_world(prediction, R=cam['orientation'], t=0) prediction[:, :, 2] -= np.min(prediction[:, :, 2]) ground_truth = camera_to_world(ground_truth, R=cam['orientation'], t=0) ground_truth[:, :, 2] -= np.min(ground_truth[:, :, 2]) anim_output = {'Regression': prediction, 'Ground truth': ground_truth} input_keypoints = image_coordinates(input_keypoints[..., :2], w=cam['res_w'], h=cam['res_h']) render_animation(input_keypoints, anim_output, dataset.skeleton(), dataset.fps(), args.viz_bitrate, cam['azimuth'], args.viz_output, limit=args.viz_limit, downsample=args.viz_downsample, size=args.viz_size, input_video_path=args.viz_video, viewport=(cam['res_w'], cam['res_h']), input_video_skip=args.viz_skip)
out_poses_2d[i] = out_poses_2d[i][::stride] if out_poses_3d is not None: out_poses_3d[i] = out_poses_3d[i][::stride] return out_camera_params, out_poses_3d, out_poses_2d action_filter = None if args.actions == '*' else args.actions.split(',') if action_filter is not None: print('Selected actions:', action_filter) cameras_valid, poses_valid, poses_valid_2d = fetch(subjects_test, action_filter) filter_widths = [int(x) for x in args.architecture.split(",")] adj = adj_mx_from_skeleton(dataset.skeleton()) if not args.disable_optimizations and not args.dense and args.stride == 1: # Use optimized model for single-frame predictions model_pos_train = SpatialTemporalModelOptimized1f( adj, poses_valid_2d[0].shape[-2], poses_valid_2d[0].shape[-1], dataset.skeleton().num_joints(), filter_widths=filter_widths, causal=args.causal, dropout=args.dropout, channels=args.channels) else: # When incompatible settings are detected (stride > 1, dense filters, or disabled optimization) fall back to normal model model_pos_train = SpatialTemporalModel(adj,
import os import argparse from tool.mpii_coco_h36m import coco_h36m, mpii_h36m from common.skeleton import Skeleton from common.graph_utils import adj_mx_from_skeleton from common.camera import normalize_screen_coordinates, camera_to_world from common.generators import * from model.gast_net import * from tool.visualization import render_animation h36m_skeleton = Skeleton( parents=[-1, 0, 1, 2, 0, 4, 5, 0, 7, 8, 9, 8, 11, 12, 8, 14, 15], joints_left=[6, 7, 8, 9, 10, 16, 17, 18, 19, 20, 21, 22, 23], joints_right=[1, 2, 3, 4, 5, 24, 25, 26, 27, 28, 29, 30, 31]) adj = adj_mx_from_skeleton(h36m_skeleton) joints_left, joints_right = [4, 5, 6, 11, 12, 13], [1, 2, 3, 14, 15, 16] kps_left, kps_right = [4, 5, 6, 11, 12, 13], [1, 2, 3, 14, 15, 16] rot = np.array([0.14070565, -0.15007018, -0.7552408, 0.62232804], dtype=np.float32) keypoints_metadata = { 'keypoints_symmetry': (joints_left, joints_right), 'layout_name': 'Human3.6M', 'num_joints': 17 } mpii_metadata = { 'layout_name': 'mpii', 'num_joints': 16, 'keypoints_symmetry': [ [3, 4, 5, 13, 14, 15],
def main(args): print('==> Using settings {}'.format(args)) print('==> Loading dataset...') if args.dataset == 'h36m': dataset_path = path.join('data', 'data_3d_' + args.dataset + '.npz') from common.h36m_dataset import Human36mDataset, TRAIN_SUBJECTS, TEST_SUBJECTS dataset = Human36mDataset(dataset_path) subjects_train = TRAIN_SUBJECTS subjects_test = TEST_SUBJECTS print('==> Preparing data ' + args.dataset + "...") dataset = read_3d_data(dataset) adj = adj_mx_from_skeleton(dataset.skeleton()) nodes_group = dataset.skeleton().joints_group() if args.non_local else None print('==> Loading 2D detections...') keypoints = create_2d_data(path.join('data', 'data_2d_' + args.dataset + '_' + args.keypoints + '.npz'), dataset) action_filter = None if args.actions == '*' else args.actions.split(',') stride = args.downsample if action_filter is not None: action_filter = map(lambda x: dataset.define_actions(x)[0], action_filter) print('==> Selected actions: {}'.format(action_filter)) if not args.evaluate: print('==> Build DataLoader...') poses_train, poses_train_2d, actions_train = fetch(subjects_train, dataset, keypoints, action_filter, stride) train_loader = DataLoader(PoseGenerator(poses_train, poses_train_2d, actions_train), batch_size=args.batch_size, shuffle=True, num_workers=args.num_workers, pin_memory=True) poses_valid, poses_valid_2d, actions_valid = fetch(subjects_test, dataset, keypoints, action_filter, stride) valid_loader = DataLoader(PoseGenerator(poses_valid, poses_valid_2d, actions_valid), batch_size=args.batch_size, shuffle=False, num_workers=args.num_workers, pin_memory=True) elif "synmit" in args.dataset: dataset_path = args.dataset_path from common.synmit_dataset import SynDataset17, SynDataset17_h36m if "h36m" in args.dataset: train_dataset = SynDataset17_h36m(dataset_path, image_set="train") valid_dataset = SynDataset17_h36m(dataset_path, image_set="val") else: train_dataset = SynDataset17(dataset_path, image_set="train") valid_dataset = SynDataset17(dataset_path, image_set="val") dataset = train_dataset adj = adj_mx_from_edges(dataset.jointcount(), dataset.edge(), sparse=False) nodes_group = dataset.nodesgroup() if args.non_local else None train_loader = DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True, num_workers=args.num_workers, pin_memory=True) valid_loader = DataLoader(valid_dataset, batch_size=args.batch_size, shuffle=False, num_workers=args.num_workers, pin_memory=True) else: raise KeyError('Invalid dataset') cudnn.benchmark = True device = torch.device("cuda") # Create model print("==> Creating model...") p_dropout = (None if args.dropout == 0.0 else args.dropout) view_count = dataset.view_count if args.multiview else None model_pos = MultiviewSemGCN(adj, args.hid_dim, coords_dim=(3, 1), num_layers=args.num_layers, p_dropout=p_dropout, view_count=view_count, nodes_group=nodes_group).to(device) print("==> Total parameters: {:.2f}M".format(sum(p.numel() for p in model_pos.parameters()) / 1000000.0)) criterion = nn.MSELoss(reduction='mean').to(device) optimizer = torch.optim.Adam(model_pos.parameters(), lr=args.lr) # Optionally resume from a checkpoint if args.resume or args.evaluate: ckpt_path = (args.resume if args.resume else args.evaluate) if path.isfile(ckpt_path): print("==> Loading checkpoint '{}'".format(ckpt_path)) ckpt = torch.load(ckpt_path) start_epoch = ckpt['epoch'] error_best = ckpt['error'] glob_step = ckpt['step'] lr_now = ckpt['lr'] # lr_now = args.lr #### if args.dataset == "h36m": for k in list(ckpt['state_dict'].keys()): v = ckpt['state_dict'].pop(k) if (type(k) == str) and ("nonlocal" in k): k = k.replace("nonlocal","nonlocal_layer") ckpt['state_dict'][k] = v model_pos.load_state_dict(ckpt['state_dict']) optimizer.load_state_dict(ckpt['optimizer']) print("==> Loaded checkpoint (Epoch: {} | Error: {})".format(start_epoch, error_best)) if args.resume: ckpt_dir_path = path.dirname(ckpt_path) logger = Logger(path.join(ckpt_dir_path, 'log.txt'), resume=True) else: raise RuntimeError("==> No checkpoint found at '{}'".format(ckpt_path)) else: start_epoch = 0 error_best = None glob_step = 0 lr_now = args.lr mv_str = "mv_" if args.multiview else "" ckpt_dir_path = path.join(args.checkpoint, args.dataset + "_" + mv_str + datetime.datetime.now().isoformat(timespec="seconds")) ckpt_dir_path = ckpt_dir_path.replace(":","-") if not path.exists(ckpt_dir_path): os.makedirs(ckpt_dir_path) print('==> Making checkpoint dir: {}'.format(ckpt_dir_path)) logger = Logger(os.path.join(ckpt_dir_path, 'log.txt')) logger.set_names(['epoch', 'lr', 'loss_train', 'error_eval_p1', 'error_eval_p2']) if args.evaluate: print('==> Evaluating...') if args.dataset == 'h36m': if action_filter is None: action_filter = dataset.define_actions() errors_p1 = np.zeros(len(action_filter)) errors_p2 = np.zeros(len(action_filter)) for i, action in enumerate(action_filter): poses_valid, poses_valid_2d, actions_valid = fetch(subjects_test, dataset, keypoints, [action], stride) valid_loader = DataLoader(PoseGenerator(poses_valid, poses_valid_2d, actions_valid), batch_size=args.batch_size, shuffle=False, num_workers=args.num_workers, pin_memory=True) errors_p1[i], errors_p2[i] = evaluate(valid_loader, model_pos, device) elif "synmit" in args.dataset: if "h36m" in args.dataset: test_dataset = SynDataset17_h36m(dataset_path, image_set="test") else: test_dataset = SynDataset17(dataset_path, image_set="test") test_loader = DataLoader(test_dataset, batch_size=args.batch_size, shuffle=False, num_workers=args.num_workers, pin_memory=True) errors_p1, errors_p2 = evaluate(test_loader, model_pos, device) print('Protocol #1 (MPJPE) action-wise average: {:.2f} (mm)'.format(np.mean(errors_p1).item())) print('Protocol #2 (REL-MPJPE) action-wise average: {:.2f} (mm)'.format(np.mean(errors_p2).item())) exit(0) epoch_loss = 1e5 for epoch in range(start_epoch, args.epochs): print('\nEpoch: %d | LR: %.8f' % (epoch + 1, lr_now)) # Train for one epoch epoch_loss, lr_now, glob_step = train(epoch_loss, train_loader, model_pos, criterion, optimizer, device, args.lr, lr_now, glob_step, args.lr_decay, args.lr_gamma, max_norm=args.max_norm, loss_3d=args.loss3d, earlyend=args.earlyend) # Evaluate error_eval_p1, error_eval_p2 = evaluate(valid_loader, model_pos, device) # Update log file logger.append([epoch + 1, lr_now, epoch_loss, error_eval_p1, error_eval_p2]) # Save checkpoint if error_best is None or error_best > error_eval_p1: error_best = error_eval_p1 save_ckpt({'epoch': epoch + 1, 'lr': lr_now, 'step': glob_step, 'state_dict': model_pos.state_dict(), 'optimizer': optimizer.state_dict(), 'error': error_eval_p1}, ckpt_dir_path, suffix='best') if (epoch + 1) % args.snapshot == 0: save_ckpt({'epoch': epoch + 1, 'lr': lr_now, 'step': glob_step, 'state_dict': model_pos.state_dict(), 'optimizer': optimizer.state_dict(), 'error': error_eval_p1}, ckpt_dir_path) logger.close() logger.plot(['loss_train', 'error_eval_p1']) savefig(path.join(ckpt_dir_path, 'log.eps')) return
def main(): dataset_path = "./data/data_3d_h36m.npz" # 加载数据 from common.h36m_dataset import Human36mDataset dataset = Human36mDataset(dataset_path) dataset = read_3d_data(dataset) cudnn.benchmark = True device = torch.device("cpu") from models.sem_gcn import SemGCN from common.graph_utils import adj_mx_from_skeleton p_dropout = None adj = adj_mx_from_skeleton(dataset.skeleton()) model_pos = SemGCN(adj, 128, num_layers=4, p_dropout=p_dropout, nodes_group=dataset.skeleton().joints_group()).to(device) ckpt_path = "./checkpoint/pretrained/ckpt_semgcn_nonlocal_sh.pth.tar" ckpt = torch.load(ckpt_path, map_location='cpu') model_pos.load_state_dict(ckpt['state_dict'], False) model_pos.eval() # ============ 新增代码 ============== # 从项目处理2d数据的代码中输出的一个人体数据 inputs_2d = [[483.0, 450], [503, 450], [503, 539], [496, 622], [469, 450], [462, 546], [469, 622], [483, 347], [483, 326], [489, 264], [448, 347], [448, 408], [441, 463], [517, 347], [524, 408], [538, 463]] # # openpose的测试样例识别结果 # inputs_2d = [[86.0, 137], [99, 128], [94, 127], [97, 110], [89, 105], [102, 129], [116, 116], [99, 110], # [105, 93], [117, 69], [147, 63], [104, 93], [89, 69], [82, 38], [89, 139], [94, 140]] inputs_2d = np.array(inputs_2d) # inputs_2d[:, 1] = np.max(inputs_2d[:, 1]) - inputs_2d[:, 1] # 变成正的人体姿态,原始数据为倒立的 cam = dataset.cameras()['S1'][0] # 获取相机参数 inputs_2d[..., :2] = normalize_screen_coordinates(inputs_2d[..., :2], w=cam['res_w'], h=cam['res_h']) # 2d坐标处理 # 画出归一化屏幕坐标并且标记序号的二维关键点图像 print(inputs_2d) # 打印归一化后2d关键点坐标 d_x = inputs_2d[:, 0] d_y = inputs_2d[:, 1] plt.figure() plt.scatter(d_x, d_y) for i, txt in enumerate(np.arange(inputs_2d.shape[0])): plt.annotate(txt, (d_x[i], d_y[i])) # 标号 # plt.show() # 显示2d关键点归一化后的图像 # 获取3d结果 inputs_2d = torch.tensor(inputs_2d, dtype=torch.float32) # 转换为张量 outputs_3d = model_pos(inputs_2d).cpu() # 加载模型 outputs_3d[:, :, :] -= outputs_3d[:, :1, :] # Remove global offset / 移除全球偏移 predictions = [outputs_3d.detach().numpy()] # 预测结果 prediction = np.concatenate(predictions)[0] # 累加取第一个 # Invert camera transformation / 反相机的转换 prediction = camera_to_world(prediction, R=cam['orientation'], t=0) # R和t的参数设置影响不大,有多种写法和选取的相机参数有关,有些S没有t等等问题 prediction[:, 2] -= np.min(prediction[:, 2]) # 向上偏移min(prediction[:, 2]),作用是把坐标变为正数 print('prediction') print(prediction) # 打印画图的3d坐标 plt.figure() ax = plt.subplot(111, projection='3d') # 创建一个三维的绘图工程 o_x = prediction[:, 0] o_y = prediction[:, 1] o_z = prediction[:, 2] print(o_x) print(o_y) print(o_z) ax.scatter(o_x, o_y, o_z) temp = o_x x = [temp[9], temp[8], temp[7], temp[10], temp[11], temp[12]] temp = o_y y = [temp[9], temp[8], temp[7], temp[10], temp[11], temp[12]] temp = o_z z = [temp[9], temp[8], temp[7], temp[10], temp[11], temp[12]] ax.plot(x, y, z) temp = o_x x = [temp[7], temp[0], temp[4], temp[5], temp[6]] temp = o_y y = [temp[7], temp[0], temp[4], temp[5], temp[6]] temp = o_z z = [temp[7], temp[0], temp[4], temp[5], temp[6]] ax.plot(x, y, z) temp = o_x x = [temp[0], temp[1], temp[2], temp[3]] temp = o_y y = [temp[0], temp[1], temp[2], temp[3]] temp = o_z z = [temp[0], temp[1], temp[2], temp[3]] ax.plot(x, y, z) temp = o_x x = [temp[7], temp[13], temp[14], temp[15]] temp = o_y y = [temp[7], temp[13], temp[14], temp[15]] temp = o_z z = [temp[7], temp[13], temp[14], temp[15]] ax.plot(x, y, z) # temp = o_x # x = [temp[0], temp[14]] # temp = o_y # y = [temp[0], temp[14]] # temp = o_z # z = [temp[0], temp[14]] # ax.plot(y, x, z) # # temp = o_x # x = [temp[0], temp[15]] # temp = o_y # y = [temp[0], temp[15]] # temp = o_z # z = [temp[0], temp[15]] # ax.plot(y, x, z) # 改变坐标比例的代码,该代码的效果是z坐标轴是其他坐标的两倍 from matplotlib.pyplot import MultipleLocatort major_locator = MultipleLocator(0.5) ax.xaxis.set_major_locator(major_locator) ax.yaxis.set_major_locator(major_locator) ax.zaxis.set_major_locator(major_locator) ax.get_proj = lambda: np.dot(Axes3D.get_proj(ax), np.diag([0.5, 0.5, 1, 1])) plt.show()
from tools.preprocess import load_kpts_json, h36m_coco_format, revise_kpts, revise_skes from tools.inference import gen_pose from tools.vis_kpts import plot_keypoint cur_dir, chk_root, data_root, lib_root, output_root = get_path(__file__) model_dir = chk_root + 'gastnet/' sys.path.insert(1, lib_root) from lib.pose import gen_video_kpts as hrnet_pose sys.path.pop(1) sys.path.pop(0) skeleton = Skeleton( parents=[-1, 0, 1, 2, 0, 4, 5, 0, 7, 8, 9, 8, 11, 12, 8, 14, 15], joints_left=[6, 7, 8, 9, 10, 16, 17, 18, 19, 20, 21, 22, 23], joints_right=[1, 2, 3, 4, 5, 24, 25, 26, 27, 28, 29, 30, 31]) adj = adj_mx_from_skeleton(skeleton) joints_left, joints_right = [4, 5, 6, 11, 12, 13], [1, 2, 3, 14, 15, 16] kps_left, kps_right = [4, 5, 6, 11, 12, 13], [1, 2, 3, 14, 15, 16] rot = np.array([0.14070565, -0.15007018, -0.7552408, 0.62232804], dtype=np.float32) keypoints_metadata = { 'keypoints_symmetry': (joints_left, joints_right), 'layout_name': 'Human3.6M', 'num_joints': 17 } width, height = (1920, 1080) def load_model_realtime(rf=81): if rf == 27:
def __init__(self, args, device, skeleton): print("Initialize SemGCMDriver - begin.") from models.sem_gcn import SemGCN from common.graph_utils import adj_mx_from_skeleton self.hid_dim = args.sem_hid_dim self.num_layers = args.sem_num_layers self.p_dropout = (None if args.sem_dropout == 0.0 else args.sem_dropout) self.render_score_threshold = args.render_score_threshold self.skeleton = skeleton adj = adj_mx_from_skeleton(self.skeleton) self.device = device self.model_pos = SemGCN(adj, self.hid_dim, num_layers=self.num_layers, p_dropout=self.p_dropout, nodes_group=self.skeleton.joints_group()).to( self.device.device) self.last_2d_positions = None self.last_3d_positions = None self.last_scores = None self.render_3d = args.sem_show_3d self._plot = args.sem_plot self._plot_initalized = False self._plot_skeleton = None self.last_3d_skeletons = None # Resume from a checkpoint ckpt_path = args.sem_evaluate if path.isfile(ckpt_path): print("==> Loading checkpoint '{}'".format(ckpt_path)) self.ckpt = torch.load(ckpt_path) self.rename_nonlocal_node(self.ckpt) start_epoch = self.ckpt['epoch'] error_best = self.ckpt['error'] self.model_pos.load_state_dict(self.ckpt['state_dict']) print("==> Loaded checkpoint (Epoch: {} | Error: {})".format( start_epoch, error_best)) else: raise RuntimeError( "==> No checkpoint found at '{}'".format(ckpt_path)) if self._plot: size = 5 radius = 1 azim = 45 title = "test" plt.ion() fig = plt.figure(figsize=(size, size)) self._plot_ax = fig.add_subplot(1, 1, 1, projection='3d') self._plot_ax.view_init(elev=15., azim=azim) self._plot_ax.set_xlim3d([-radius / 2, radius / 2]) self._plot_ax.set_zlim3d([0, radius]) self._plot_ax.set_ylim3d([-radius / 2, radius / 2]) self._plot_ax.set_aspect('auto') self._plot_ax.set_xticklabels([]) self._plot_ax.set_yticklabels([]) self._plot_ax.set_zticklabels([]) self._plot_ax.dist = 7.5 self._plot_ax.set_title(title) # , pad=35 print("Initialize SemGCMDriver - end.")
def reconstruction(args): """ Generate 3D poses from 2D keypoints detected from video, and visualize it :param chk_file: The file path of model weight :param kps_file: The file path of 2D keypoints :param viz_output: The output path of animation :param video_path: The input video path :param kpts_format: The format of 2D keypoints, like MSCOCO, MPII, H36M, OpenPose. The default format is H36M """ # Getting joint information joints_left, joints_right, h36m_skeleton, keypoints_metadata = get_joints_info( args.num_joints) kps_left, kps_right = joints_left, joints_right adj = adj_mx_from_skeleton(h36m_skeleton) print('Loading 2D keypoints ...') keypoints, scores, _, _ = load_json(args.keypoints_file, args.num_joints) # Loading only one person's keypoints if len(keypoints.shape) == 4: keypoints = keypoints[0] assert len(keypoints.shape) == 3 # Transform the keypoints format from different dataset (MSCOCO, MPII) to h36m format if args.kpts_format == 'coco': keypoints, valid_frames = coco_h36m(keypoints) elif args.kpts_format == 'mpii': keypoints, valid_frames = mpii_h36m(keypoints) elif args.kpts_format == 'openpose': # Convert 'Openpose' format to MSCOCO order_coco = [i for i in range(18) if i != 1] keypoints = keypoints[:, order_coco] keypoints, valid_frames = coco_h36m(keypoints) elif args.kpts_format == 'wholebody': keypoints, valid_frames = coco_h36m_toe_format(keypoints) else: valid_frames = np.where( np.sum(keypoints.reshape(-1, 34), axis=1) != 0)[0] assert args.kpts_format == 'h36m' # Get the width and height of video cap = cv2.VideoCapture(args.video_path) width = int(round(cap.get(cv2.CAP_PROP_FRAME_WIDTH))) height = int(round(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))) # normalize keypoints input_keypoints = normalize_screen_coordinates(keypoints[..., :2], w=width, h=height) if args.frames == 27: filter_widths = [3, 3, 3] channels = 128 elif args.frames == 81: filter_widths = [3, 3, 3, 3] channels = 64 else: filter_widths = [3, 3, 3, 3, 3] channels = 32 model_pos = SpatioTemporalModel(adj, args.num_joints, 2, args.num_joints, filter_widths=filter_widths, channels=channels, dropout=0.05, causal=args.causal) if torch.cuda.is_available(): model_pos = model_pos.cuda() # load pretrained model print('Loading checkpoint', args.weight) chk_file = os.path.join('./checkpoint/gastnet', args.weight) checkpoint = torch.load(chk_file, map_location=lambda storage, loc: storage) model_pos.load_state_dict(checkpoint['model_pos']) receptive_field = model_pos.receptive_field() pad = (receptive_field - 1) // 2 # Padding on each side if args.causal: causal_shift = pad else: causal_shift = 0 print('Reconstructing ...') gen = UnchunkedGenerator(None, None, [input_keypoints[valid_frames]], pad=pad, causal_shift=causal_shift, augment=True, kps_left=kps_left, kps_right=kps_right, joints_left=joints_left, joints_right=joints_right) prediction = evaluate(gen, model_pos, joints_left, joints_right, return_predictions=True) prediction = camera_to_world(prediction, R=rot, t=0) # We don't have the trajectory, but at least we can rebase the height prediction[:, :, 2] -= np.min(prediction[:, :, 2]) prediction_new = np.zeros((*input_keypoints.shape[:-1], 3), dtype=np.float32) prediction_new[valid_frames] = prediction print('Rendering ...') anim_output = {'Reconstruction': prediction_new} render_animation(keypoints, keypoints_metadata, anim_output, h36m_skeleton, 25, 3000, np.array(70., dtype=np.float32), args.viz_output, limit=-1, downsample=1, size=5, input_video_path=args.video_path, viewport=(width, height), input_video_skip=0)