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)
