def main():
args = get_args()
printer(args)
# Initialize logger
global log
log = Logger(args)
# Set random seed
random.seed(args.seed)
torch.manual_seed(args.seed)
if args.cuda:
torch.cuda.manual_seed(args.seed)
# Check cuda device
device = torch.device('cuda' if args.cuda else 'cpu')
# Loading data
train_music_data, train_dance_data, _ = load_data(
args.train_dir, interval=args.seq_len, data_type=args.data_type)
training_data = prepare_dataloader(train_music_data, train_dance_data, args)
valid_music_data, valid_dance_data, _ = load_data(
args.valid_dir, interval=args.seq_len, data_type=args.data_type)
validation_data = prepare_dataloader(valid_music_data, valid_dance_data, args)
encoder = Encoder(args)
decoder = Decoder(args)
model = Model(encoder, decoder, args, device=device)
for name, parameters in model.named_parameters():
print(name, ':', parameters.size())
# Data Parallel to use multi-gpu
model = nn.DataParallel(model).to(device)
#model = model.to(device)
optimizer = optim.Adam(filter(lambda x: x.requires_grad, model.module.parameters()), lr=args.lr)
# scheduler = ReduceLROnPlateau(
# optimizer, mode='min', factor=0.8, patience=30, verbose=True, min_lr=1e-06, eps=1e-07)
# scheduler = optim.lr_scheduler.StepLR(optimizer=optimizer,
# step_size=args.lr_step_size,
# gamma=0.5)
# for name,para in model.named_parameters():
# print(name,para.size())
train(model, training_data, validation_data, optimizer, device, args, log)
def main():
args = get_args()
printer(args)
# Initialize logger
global log
log = Logger(args)
# Set random seed
random.seed(args.seed)
torch.manual_seed(args.seed)
if args.cuda:
torch.cuda.manual_seed(args.seed)
# Check cuda device
device = torch.device('cuda' if args.cuda else 'cpu')
# Loading data
if args.aist:
print ("train with AIST++ dataset!")
train_music_data, train_dance_data, _ = load_data_aist(
args.train_dir, interval=args.seq_len, rotmat=args.rotmat)
else:
train_music_data, train_dance_data, _ = load_data(
args.train_dir, interval=args.seq_len, data_type=args.data_type)
training_data = prepare_dataloader(train_music_data, train_dance_data, args)
print ("data shape:", len(training_data))
encoder = Encoder(args)
decoder = Decoder(args)
model = Model(encoder, decoder, args, device=device)
# for name, parameters in model.named_parameters():
# print(name, ':', parameters.size())
# Data Parallel to use multi-gpu
model = nn.DataParallel(model).to(device)
#model = model.to(device)
optimizer = optim.Adam(filter(lambda x: x.requires_grad, model.module.parameters()), lr=args.lr)
train(model, training_data, optimizer, device, args, log)
def main():
args = get_args()
if not os.path.exists(args.json_dir):
os.makedirs(args.json_dir)
music_data, dance_data, dance_names = load_data(args.input_dir,
interval=None,
data_type=args.data_type)
device = torch.device('cuda' if args.cuda else 'cpu')
test_loader = torch.utils.data.DataLoader(DanceDataset(
music_data, dance_data),
batch_size=args.batch_size,
collate_fn=paired_collate_fn)
generator = Generator(args.model, device)
results = []
for batch in tqdm(test_loader, desc='Generating dance poses'):
# Prepare data
src_seq, src_pos, _ = map(lambda x: x.to(device), batch)
pose_seq = generator.generate(src_seq, src_pos)
results.append(pose_seq)
# Visualize generated dance poses
np_dances = []
for i in range(len(results)):
np_dance = results[i][0].data.cpu().numpy()
root = np_dance[:, 2 * 8:2 * 9]
np_dance = np_dance + np.tile(root, (1, 25))
np_dance[:, 2 * 8:2 * 9] = root
np_dances.append(np_dance)
write2json(np_dances, dance_names, args)
visualize(args)
def main():
""" Main function """
parser = argparse.ArgumentParser()
parser.add_argument('--train_dir',
type=str,
default='data/train_1min',
help='the directory of dance data')
parser.add_argument('--test_dir',
type=str,
default='data/test_1min',
help='the directory of music feature data')
parser.add_argument('--data_type',
type=str,
default='2D',
help='the type of training data')
parser.add_argument(
'--output_dir',
metavar='PATH',
default='checkpoints/layers2_win100_schedule100_condition10_detach')
parser.add_argument('--epoch', type=int, default=10000)
parser.add_argument('--batch_size', type=int, default=32)
parser.add_argument('--save_per_epochs',
type=int,
metavar='N',
default=500)
parser.add_argument('--log_per_updates',
type=int,
metavar='N',
default=1,
help='log model loss per x updates (mini-batches).')
parser.add_argument('--seed',
type=int,
default=1234,
help='random seed for data shuffling, dropout, etc.')
parser.add_argument('--tensorboard', action='store_false')
parser.add_argument('--d_frame_vec', type=int, default=438)
parser.add_argument('--frame_emb_size', type=int, default=200)
parser.add_argument('--d_pose_vec', type=int, default=50)
parser.add_argument('--pose_emb_size', type=int, default=50)
parser.add_argument('--d_inner', type=int, default=1024)
parser.add_argument('--d_k', type=int, default=64)
parser.add_argument('--d_v', type=int, default=64)
parser.add_argument('--n_head', type=int, default=8)
parser.add_argument('--n_layers', type=int, default=2)
parser.add_argument('--lr', type=float, default=1e-4)
parser.add_argument('--dropout', type=float, default=0.1)
parser.add_argument('--seq_len', type=int, default=900)
parser.add_argument('--max_seq_len', type=int, default=4500)
parser.add_argument('--condition_step', type=int, default=10)
parser.add_argument('--sliding_windown_size', type=int, default=100)
parser.add_argument('--lambda_v', type=float, default=0.01)
parser.add_argument('--cuda',
type=str2bool,
nargs='?',
metavar='BOOL',
const=True,
default=torch.cuda.is_available(),
help='whether to use GPU acceleration.')
args = parser.parse_args()
args.d_model = args.frame_emb_size
if not os.path.exists(args.output_dir):
os.makedirs(args.output_dir)
global log
log = Logger(args)
print(args)
# Set random seed
random.seed(args.seed)
torch.manual_seed(args.seed)
if args.cuda:
torch.cuda.manual_seed(args.seed)
# Loading training data
train_music_data, train_dance_data = load_data(args.train_dir,
interval=args.seq_len,
data_type=args.data_type)
training_data = prepare_dataloader(train_music_data, train_dance_data,
args)
device = torch.device('cuda' if args.cuda else 'cpu')
encoder = Encoder(max_seq_len=args.max_seq_len,
input_size=args.d_frame_vec,
d_word_vec=args.frame_emb_size,
n_layers=args.n_layers,
n_head=args.n_head,
d_k=args.d_k,
d_v=args.d_v,
d_model=args.d_model,
d_inner=args.d_inner,
dropout=args.dropout)
decoder = Decoder(input_size=args.d_pose_vec,
d_word_vec=args.pose_emb_size,
hidden_size=args.d_inner,
encoder_d_model=args.d_model,
dropout=args.dropout)
model = Model(encoder,
decoder,
condition_step=args.condition_step,
sliding_windown_size=args.sliding_windown_size,
lambda_v=args.lambda_v,
device=device)
print(model)
for name, parameters in model.named_parameters():
print(name, ':', parameters.size())
# Data Parallel to use multi-gpu
model = nn.DataParallel(model).to(device)
optimizer = optim.Adam(filter(lambda x: x.requires_grad,
model.module.parameters()),
lr=args.lr)
train(model, training_data, optimizer, device, args, log)
def main():
args = get_args()
if args.aist:
import vedo
if not os.path.exists(args.json_dir):
os.makedirs(args.json_dir)
if args.aist:
print ("test with AIST++ dataset!")
music_data, dance_data, dance_names = load_data_aist(
args.input_dir, interval=None, rotmat=args.rotmat)
else:
music_data, dance_data, dance_names = load_data(
args.input_dir, interval=None)
device = torch.device('cuda' if args.cuda else 'cpu')
test_loader = torch.utils.data.DataLoader(
DanceDataset(music_data, dance_data),
batch_size=args.batch_size,
collate_fn=paired_collate_fn
)
generator = Generator(args.model, device)
if args.aist and args.rotmat:
from smplx import SMPL
smpl = SMPL(model_path="/media/ruilongli/hd1/Data/smpl/", gender='MALE', batch_size=1)
results = []
random_id = 0 # np.random.randint(0, 1e4)
for i, batch in enumerate(tqdm(test_loader, desc='Generating dance poses')):
# Prepare data
src_seq, src_pos, tgt_pose = map(lambda x: x.to(device), batch)
pose_seq = generator.generate(src_seq[:, :1200], src_pos[:, :1200]) # first 20 secs
results.append(pose_seq)
if args.aist:
np_dance = pose_seq[0].data.cpu().numpy()
if args.rotmat:
root = np_dance[:, :3]
rotmat = np_dance[:, 3:].reshape([-1, 3, 3])
rotmat = get_closest_rotmat(rotmat)
smpl_poses = rotmat2aa(rotmat).reshape(-1, 24, 3)
np_dance = smpl.forward(
global_orient=torch.from_numpy(smpl_poses[:, 0:1]).float(),
body_pose=torch.from_numpy(smpl_poses[:, 1:]).float(),
transl=torch.from_numpy(root).float(),
).joints.detach().numpy()[:, 0:24, :]
else:
root = np_dance[:, :3]
np_dance = np_dance + np.tile(root, (1, 24))
np_dance[:, :3] = root
np_dance = np_dance.reshape(np_dance.shape[0], -1, 3)
print (np_dance.shape)
# save
save_path = os.path.join(args.json_dir, dance_names[i]+f"_{random_id:04d}")
np.save(save_path, np_dance)
# visualize
for frame in np_dance:
pts = vedo.Points(frame, r=20)
vedo.show(pts, interactive=False)
time.sleep(0.1)
exit()
if args.aist:
pass
else:
# Visualize generated dance poses
np_dances = []
for i in range(len(results)):
np_dance = results[i][0].data.cpu().numpy()
root = np_dance[:, 2*8:2*9]
np_dance = np_dance + np.tile(root, (1, 25))
np_dance[:, 2*8:2*9] = root
np_dances.append(np_dance)
write2json(np_dances, dance_names, args)
visualize(args)