def test(): args = parse_args() cfg = Config.from_file(args.config) data_path = cfg.test.source_data out = cfg.test.out model_path = cfg.test.gen_path print('GPU: {}'.format(args.gpu)) # number of gesture classes and users n_gesture = cfg.train.n_gesture n_user = len(cfg.train.dataset_dirs) n_style = n_gesture if cfg.style == 'gesture' else n_user ## Import generator model gen = getattr(models, cfg.train.generator.model)(cfg.train.generator, n_style=n_style) serializers.load_npz(model_path, gen) print('') print(f'loading generator weight from {model_path}') ## Set GPU if args.gpu >= 0: # Make a specified GPU current chainer.backends.cuda.get_device_from_id(args.gpu).use() gen.to_gpu() test_data = data_load(data_path, ges_class=cfg.test.ges_class) test_data = np.expand_dims(test_data, axis=1) style_label = np.zeros((test_data.shape[0], n_style * 2, 1, 1)) source = cfg.test.ges_class if cfg.style == 'gesture' else cfg.source_user target = cfg.test.target_style style_label[:, source] += 1 style_label[:, target + n_style] += 1 test_data = Variable(cuda.to_gpu(test_data.astype(np.float32))) style_label = Variable(cuda.to_gpu(style_label.astype(np.float32))) with chainer.using_config('train', False), chainer.using_config( 'enable_backprop', False): gen_data = gen(test_data, style_label) gen_data.to_cpu() if cfg.style == 'gesture': save_path = f'./{out}/user{cfg.source_user}/ges{target}_from_ges{source}' else: save_path = f'./{out}/user{target}/ges{cfg.test.ges_class}_from_user{source}' if not os.path.exists(os.path.dirname(save_path)): os.makedirs(os.path.dirname(save_path)) print('saving generated data to ' + save_path + '.npy') np.save(save_path, gen_data.data)
def main(): checkpoint_path = '/DATA/wangshen_data/UCF101/voxelflow_finetune_model_best.pth.tar' global cfg, best_PSNR args = parse_args() cfg = Config.from_file(args.config) str1 = ','.join(str(gpu) for gpu in cfg.device) print(str1) os.environ["CUDA_VISIBLE_DEVICES"] = str1 # os.environ["CUDA_VISIBLE_DEVICES"] = '0' cudnn.benchmark = True cudnn.fastest = True if hasattr(datasets, cfg.dataset): ds = getattr(datasets, cfg.dataset) else: raise ValueError('Unknown dataset ' + cfg.dataset) model = getattr(models, cfg.model.name)(cfg.model).cuda() cfg.train.input_mean = model.input_mean cfg.train.input_std = model.input_std cfg.test.input_mean = model.input_mean cfg.test.input_std = model.input_std input_mean = cfg.test.input_mean input_std = cfg.test.input_std # Data loading code val_loader = torch.utils.data.DataLoader( datasets.VimeoTest(cfg.test), batch_size=1, shuffle=False, num_workers=0, #32, pin_memory=True) if os.path.isfile(checkpoint_path): print(("=> loading checkpoint '{}'".format(checkpoint_path))) checkpoint = torch.load(checkpoint_path) model.load_state_dict(checkpoint['state_dict'], False) else: print(("=> no checkpoint found at '{}'".format(checkpoint_path))) model = DataParallelwithSyncBN( model, device_ids=range(len(cfg.device))).cuda() # define loss function (criterion) optimizer and evaluator criterion = torch.nn.MSELoss().cuda() evaluator = EvalPSNR(255.0 / np.mean(cfg.test.input_std)) PSNR = validate(val_loader, model, criterion, evaluator, input_mean, input_std) print(PSNR)
def calcurate_average_speed(): parser = argparse.ArgumentParser() parser.add_argument('config') parser.add_argument('--out', required=True) args = parser.parse_args() cfg = Config.from_file(args.config) npy_paths = collect_path(cfg.train.dataset) head, ext = os.path.splitext(cfg.train.dataset) head, data_name = os.path.split(head) class_list = cfg.train.class_list v_traj_avg_dict = {label: 0 for label in class_list} class_npy_paths = {label: [] for label in class_list} for npy_path in npy_paths: npy_name = os.path.splitext(os.path.split(npy_path)[1])[0] label = npy_name.split('_')[0] class_npy_paths[label].append(npy_path) for label in class_npy_paths.keys(): class_v_traj = None for npy_path in class_npy_paths[label]: motion = np.load(npy_path) trajX = motion[:, 0] trajZ = motion[:, 2] v_trajX = trajX[:-1] - trajX[1:] v_trajZ = trajZ[:-1] - trajZ[1:] class_v_traj = np.sqrt( v_trajX**2 + v_trajZ**2) if class_v_traj is None else np.concatenate( (class_v_traj, np.sqrt(v_trajX**2 + v_trajZ**2)), axis=0) if not class_npy_paths[label]: continue class_v_traj_avg = np.average(class_v_traj, axis=0) v_traj_avg_dict[label] = class_v_traj_avg print(v_traj_avg_dict) with open(args.out, 'wb') as f: pickle.dump(v_traj_avg_dict, f) return
def main(): global cfg, best_PSNR args = parse_args() cfg = Config.from_file(args.config) #os.environ["CUDA_VISIBLE_DEVICES"] = ','.join(str(gpu) for gpu in cfg.device) cudnn.benchmark = True cudnn.fastest = True #if hasattr(datasets, cfg.dataset): # ds = getattr(datasets, cfg.dataset) #else: # raise ValueError('Unknown dataset ' + cfg.dataset) model = getattr(models, cfg.model.name)(cfg.model).cuda() cfg.train.input_mean = model.input_mean cfg.train.input_std = model.input_std cfg.test.input_mean = model.input_mean cfg.test.input_std = model.input_std # Data loading code train_loader = torch.utils.data.DataLoader( #datasets.Combined240('data/combined', cfg.train), datasets.VimeoSeptuplet('data/vimeo_septuplet', cfg.train, True), # ds(cfg.train), batch_size=cfg.train.batch_size, shuffle=True, num_workers=8, pin_memory=True, drop_last=True) val_loader = torch.utils.data.DataLoader( # datasets.UCF101(cfg.test, False), datasets.VimeoSeptuplet('data/vimeo_septuplet', cfg.test, False), batch_size=cfg.test.batch_size, shuffle=False, num_workers=8, pin_memory=True) '''test_db_list = {'ucf101': 'UCF101Voxelflow', 'middlebury': 'Middlebury'} test_db = 'middlebury' test_loader = torch.utils.data.DataLoader( datasets.UCF101Voxelflow('data/UCF-101', cfg.test), # datasets.Middlebury('data/Middlebury', cfg.test), #getattr(datasets, test_db_list[test_db])('data/Middlebury', cfg.test), batch_size=1,#cfg.test.batch_size, shuffle=False, num_workers=16, pin_memory=True) ''' cfg.train.optimizer.args.max_iter = ( cfg.train.optimizer.args.max_epoch * len(train_loader)) policies = model.get_optim_policies() for group in policies: print(('group: {} has {} params, lr_mult: {}, decay_mult: {}'.format( group['name'], len(group['params']), group['lr_mult'], group['decay_mult']))) optimizer = Optim(policies, cfg.train.optimizer) if cfg.resume or cfg.weight: checkpoint_path = cfg.resume if cfg.resume else cfg.weight if os.path.isfile(checkpoint_path): print(("=> loading checkpoint '{}'".format(checkpoint_path))) checkpoint = torch.load(checkpoint_path) model.load_state_dict(checkpoint['state_dict'], False) if cfg.resume: optimizer.load_state_dict(checkpoint['grad_dict']) del checkpoint else: print(("=> no checkpoint found at '{}'".format(checkpoint_path))) #model = DataParallelwithSyncBN( # model, device_ids=range(len(cfg.device))).cuda() model = model.cuda() # define loss function (criterion) optimizer and evaluator criterion = torch.nn.MSELoss().cuda() # evaluator = EvalPSNR(255.0 / np.mean(cfg.test.input_std)) evaluator = EvalPSNR(255.0 / np.mean(cfg.test.input_std), cfg.test.input_mean, cfg.test.input_std) #PSNR = validate(val_loader, model, optimizer, criterion, evaluator) ##PSNR = test(test_loader, model, optimizer, criterion, evaluator) if cfg.mode == 'test': PSNR = validate(val_loader, model, optimizer, criterion, evaluator) return for epoch in range(cfg.train.optimizer.args.max_epoch): # train for one epoch train(train_loader, model, optimizer, criterion, epoch) # evaluate on validation set if ((epoch + 1) % cfg.logging.eval_freq == 0 or epoch == cfg.train.optimizer.args.max_epoch - 1): PSNR = validate(val_loader, model, optimizer, criterion, evaluator) # remember best PSNR and save checkpoint is_best = PSNR > best_PSNR best_PSNR = max(PSNR, best_PSNR) save_checkpoint({ 'epoch': epoch + 1, 'arch': dict(cfg), 'state_dict': model.state_dict(), 'grad_dict': optimizer.state_dict(), 'best_PSNR': best_PSNR, }, is_best)
def train(): args = parse_args() cfg = Config.from_file(args.config) iteration = cfg.train.iterations batchsize = cfg.train.batchsize out = cfg.train.out print('GPU: {}'.format(args.gpu)) print('# Minibatch-size: {}'.format(batchsize)) print('# iteration: {}'.format(iteration)) print('') ## Set up networks to train # number of gesture classes and users n_gesture = cfg.train.n_gesture n_user = len(cfg.train.dataset_dirs) if cfg.style == 'gesture': gen = getattr(models, cfg.train.generator.model)(cfg.train.generator, n_style=n_gesture) elif cfg.style == 'user': gen = getattr(models, cfg.train.generator.model)(cfg.train.generator, n_style=n_user) else: print(f'Invalid style: {cfg.style}') exit() dis = getattr(models, cfg.train.discriminator.model)(cfg.train.discriminator, n_gesture=n_gesture, n_user=n_user) ## Load resume checkpoint to restart. Load optimizer state later. if args.resume or cfg.resume: gen_resume = args.resume if args.resume else cfg.resume print(f'loading generator resume from {os.path.join(out, gen_resume)}') serializers.load_npz(os.path.join(out, gen_resume), gen) dis_resume = gen_resume.replace('gen', 'dis') print( f'loading discriminator resume from {os.path.join(out, dis_resume)}' ) serializers.load_npz(os.path.join(out, dis_resume), dis) ## Load resume checkpoint (only weight, not load optimizer state) if cfg.weight: gen_weight = cfg.weight print(f'loading generator weight from {gen_weight}') serializers.load_npz(gen_weight, gen) dis_weight = gen_weight.replace('gen', 'dis') print(f'loading discriminator weight from {dis_weight}') serializers.load_npz(dis_weight, dis) ## Setup an optimizer def make_optimizer(model, alpha=0.0002, beta1=0.5): optimizer = chainer.optimizers.Adam(alpha=alpha, beta1=beta1) optimizer.setup(model) return optimizer opt_gen = make_optimizer(gen, alpha=cfg.train.parameters.g_lr, beta1=0.5) opt_dis = make_optimizer(dis, alpha=cfg.train.parameters.d_lr, beta1=0.5) if args.resume or cfg.resume: opt_gen_resume = gen_resume.replace('gen', 'opt_gen') print( f'loading generator optimizer from {os.path.join(out, opt_gen_resume)}' ) serializers.load_npz(opt_gen_resume, opt_gen) opt_dis_resume = gen_resume.replace('gen', 'opt_dis') print( f'loading discriminator optimizer from {os.path.join(out, opt_dis_resume)}' ) serializers.load_npz(opt_dis_resume, opt_dis) ## Set GPU if args.gpu >= 0: # Make a specified GPU current chainer.backends.cuda.get_device_from_id(args.gpu).use() gen.to_gpu() dis.to_gpu() ## Set up dataset train = GestureDataset(cfg.train.dataset_dirs, style=cfg.style, equal=cfg.train.class_equal) for i in range(len(cfg.train.dataset_dirs)): print( f'{cfg.train.dataset_dirs[i]} contains {train.len_each()[i]} samples' ) train_iter = chainer.iterators.SerialIterator(train, batchsize) ## Set up a Trainer updater = StarGANUpdater(models=(gen, dis), iterator=train_iter, optimizer={ 'gen': opt_gen, 'dis': opt_dis }, cfg=cfg, out=out) trainer = training.Trainer(updater, (iteration, 'iteration'), out=out) ## Set invervals ## - display_interval : Interval iterations of print logs on display. ## - save_interval : Interval iterations of save models. display_interval = (cfg.train.display_interval, 'iteration') save_interval = (cfg.train.save_interval, 'iteration') trainer.extend(extensions.snapshot_object( gen, 'gen_iter_{.updater.iteration}.npz'), trigger=save_interval) trainer.extend(extensions.snapshot_object( dis, 'dis_iter_{.updater.iteration}.npz'), trigger=save_interval) trainer.extend(extensions.snapshot_object( opt_gen, 'opt_gen_iter_{.updater.iteration}.npz'), trigger=save_interval) trainer.extend(extensions.snapshot_object( opt_dis, 'opt_dis_iter_{.updater.iteration}.npz'), trigger=save_interval) trainer.extend(extensions.LogReport(trigger=display_interval)) trainer.extend(extensions.PrintReport([ 'epoch', 'iteration', 'lr', 'gen/loss_adv', 'gen/loss_eq', 'gen/loss_style', 'gen/loss_cont', 'gen/loss_rec', 'gen/loss_sm', 'dis/loss_adv', 'dis/loss_style', 'dis/loss_cont' ]), trigger=display_interval) trainer.extend(extensions.ProgressBar(update_interval=display_interval[0])) ## Save scripts and command to result directory. (To look back later.) if not os.path.exists(out): os.makedirs(out) shutil.copy(args.config, f'./{out}') shutil.copy('./core/models/StarGAN.py', f'./{out}') shutil.copy('./core/updater/StarGANupdater.py', f'./{out}') commands = sys.argv with open(f'./{out}/command.txt', 'w') as f: f.write(f'python {commands[0]} ') for command in commands[1:]: f.write(command + ' ') ## Run the training trainer.run() ## Once training finishid, save all models. modelname = f'./{out}/gen.npz' print('saving generator model to ' + modelname) serializers.save_npz(modelname, gen) modelname = f'./{out}/dis.npz' print('saving discriminator model to ' + modelname) serializers.save_npz(modelname, dis) optname = f'./{out}/opt_gen.npz' print('saving generator optimizer to ' + optname) serializers.save_npz(optname, opt_gen) optname = f'./{out}/opt_dis.npz' print('saving discriminator optimizer to ' + optname) serializers.save_npz(optname, opt_dis)
def main(): def save_frame(th, do_pause, dir_name='', vis=True): result = prob_to_label(combine_prob(pred_prob[th])) result_show = np.dstack( (colors[result, 0], colors[result, 1], colors[result, 2])).astype(np.uint8) if args.output != '' and dir_name != '': out_file = os.path.join(dataset_dir, 'Results', 'Segmentations', resolution, args.output, dir_name, video_dir, '%05d.png' % th) if not os.path.exists(os.path.split(out_file)[0]): os.makedirs(os.path.split(out_file)[0]) if vis: cv2.imwrite(out_file, result_show) else: cv2.imwrite(out_file, result) temp = cv2.resize(frames[th], frame_0.shape[-2::-1]) * 0.3 + result_show * 0.7 return cv2.imshow('Result', temp.astype(np.uint8)) if do_pause: cv2.waitKey() else: cv2.waitKey(100) colors = labelcolormap(256) global pred_prob, frames, flow1, flow2, orig_mask, \ model, instance_num, fr_h_r, fr_w_r, appear, bbox_cnt, \ location, patch_shapes args = parse_args() cfg = Config.from_file(args.config) os.environ["CUDA_VISIBLE_DEVICES"] = str(args.gpu) model = [] model = getattr(models, cfg.model.name)(cfg.model) if os.path.isfile(cfg.weight): print(("=> loading checkpoint '{}'".format(cfg.weight))) checkpoint = torch.load(cfg.weight) model.load_state_dict(checkpoint['state_dict'], True) print(("=> loaded checkpoint")) else: raise (("=> no checkpoint found at '{}'".format(cfg.weight))) model = model.cuda() model.eval() cudnn.benchmark = True # Setup dataset dataset_dir = os.path.join('data/DAVIS') resolution = '480p' imageset_dir = os.path.join(dataset_dir, 'ImageSets', '2017', args.testset + '.txt') video_list = [] for line in open(imageset_dir).readlines(): if line.strip() != '': video_list.append(line.strip()) person_all = pickle_load(os.path.join(dataset_dir, 'PersonSearch', resolution, args.testset + '.pkl'), encoding='latin') object_all = pickle_load(os.path.join(dataset_dir, 'ObjectSearch', resolution, args.testset + '.pkl'), encoding='latin') category_all = pickle_load(os.path.join(dataset_dir, 'Class', resolution, args.testset + '.pkl'), encoding='latin') frame_cnt = 0 use_cache = False #(args.cache != '') video_cnt = -1 for video_dir in video_list: video_cnt += 1 frame_dir = os.path.join(dataset_dir, 'JPEGImages', resolution, video_dir) frame_fr_dir = os.path.join(dataset_dir, 'JPEGImages', 'Full-Resolution', video_dir) label_dir = os.path.join(dataset_dir, 'Annotations', resolution, video_dir) flow_dir = os.path.join(dataset_dir, 'Flow', resolution, video_dir) cache_dir = os.path.join(dataset_dir, 'Cache', resolution, args.cache, video_dir) frames_num = len(os.listdir(frame_dir)) if (video_cnt % args.gpu_num != args.gpu): frame_cnt += frames_num continue frame_0 = cv2.imread(os.path.join(frame_dir, '%05d.jpg' % 0)) #label_0 = cv2.imread(os.path.join(label_dir, '%05d.png' % 0), cv2.IMREAD_UNCHANGED) label_0 = cv2.imread(os.path.join(label_dir, '%05d.png' % 0), cv2.IMREAD_GRAYSCALE) label_0 = convertColorMask2Ordinary(label_0) instance_num = label_0.max() frames = [None for _ in range(frames_num)] pred_prob = [None for _ in range(frames_num)] flow1 = [None for _ in range(frames_num)] flow2 = [None for _ in range(frames_num)] person_reid = [[None for _ in range(instance_num)] for _ in range(frames_num)] object_reid = [[None for _ in range(instance_num)] for _ in range(frames_num)] category = category_all[video_dir] orig_mask = [None for _ in range(instance_num)] frames[0] = cv2.imread(os.path.join(frame_fr_dir, '%05d.jpg' % 0)) fr_h_r = float(frames[0].shape[0]) / float(frame_0.shape[0]) fr_w_r = float(frames[0].shape[1]) / float(frame_0.shape[1]) pred_prob[0] = label_to_prob(label_0, instance_num) person_reid[0] = person_all[frame_cnt] object_reid[0] = object_all[frame_cnt] save_frame(0, False, 'result', False) bbox = gen_bbox(label_0, range(instance_num), True) for i in range(instance_num): orig_mask[i] = pred_prob[0][bbox[i, 1]:bbox[i, 3], bbox[i, 0]:bbox[i, 2], i * 2 + 1] for th in range(1, frames_num): frames[th] = cv2.imread(os.path.join(frame_fr_dir, '%05d.jpg' % th)) pred_prob[th] = label_to_prob(np.zeros_like(label_0, np.uint8), instance_num) flow1[th - 1] = flo.readFlow( os.path.join(flow_dir, '%05d.flo' % (th - 1))) flow2[th] = flo.readFlow(os.path.join(flow_dir, '%05d.rflo' % th)) person_reid[th] = person_all[frame_cnt + th] object_reid[th] = object_all[frame_cnt + th] bbox_cnt = -1000 * np.ones((frames_num, instance_num)) bbox_cnt[0, :] = 0 for th in range(frames_num): for i in range(instance_num): person_reid[th][i] = person_reid[th][i][:, [0, 1, 2, 3, 5]] object_reid[th][i] = object_reid[th][i][:, [0, 1, 2, 3, 5]] frame_cnt += frames_num cache_file = os.path.join(cache_dir, '%s.pkl' % video_dir) if (use_cache and os.path.exists(cache_file)): pred_prob, bbox_cnt = pickle_load(cache_file, encoding='latin') else: predict(1, frames_num, 1, range(instance_num)) if use_cache: if not os.path.exists(os.path.split(cache_file)[0]): os.makedirs(os.path.split(cache_file)[0]) pickle_dump((pred_prob, bbox_cnt), cache_file) appear = np.zeros((frames_num, instance_num)).astype(int) location = np.zeros((frames_num, instance_num, 2)).astype(int) update_appear() for th in range(frames_num): save_frame(th, False, 'draft', True) for reid_target in ['person', 'object']: cache_file = os.path.join(cache_dir, '%s_%s.pkl' % (reid_target, video_dir)) if (use_cache and os.path.exists(cache_file)): pred_prob, bbox_cnt = pickle_load(cache_file, encoding='latin') else: target_instance = [] for i in range(instance_num): if (reid_target == 'object' or category[i][123] > 0.5): # person is 123 target_instance.append(i) reid_score = person_reid if reid_target == 'person' else object_reid draft_cnt = 0 while (True): max_score = 0 for i in range(1, frames_num - 1): temp_label = prob_to_label(combine_prob(pred_prob[i])) bbox_i = gen_bbox(temp_label, range(instance_num), False, 0.99) for j in target_instance: if bbox_cnt[ i, j] != i and reid_score[i][j].shape[0] > 0: bbox_id = np.argmax(reid_score[i][j][:, 4]) # retrieval if (appear[i, j] == 0): x1, y1, x2, y2 = reid_score[i][j][bbox_id, 0:4] if (reid_score[i][j][bbox_id, 4] > max_score and reid_score[i][j][bbox_id, 4] > reid_th): bbox_now = reid_score[i][j][ bbox_id, 0:4].astype(int) result = np.bincount(temp_label[ bbox_now[1]:bbox_now[3] + 1, bbox_now[0]:bbox_now[2] + 1].flatten(), minlength=j + 2) flag = True if flag: for occ_instance in np.where( result[1:] > 0)[0]: if (IoU( bbox_now, bbox_i[occ_instance]) > bbox_occ_th): flag = False if flag: for k in target_instance: if (k != j and appear[i, k] == 0 and reid_score[i][k][ bbox_id, 4] > reid_score[i][j][ bbox_id, 4]): flag = False if flag: max_frame = i max_instance = j max_bbox = reid_score[i][j][ bbox_id, 0:4] max_score = reid_score[i][j][ bbox_id, 4] if (max_score == 0): break bbox_cnt[max_frame, max_instance] = max_frame predict_single(max_frame, max_instance, max_bbox, orig_mask[max_instance]) save_frame(max_frame, False, '%s_%05d_checkpoint' % (reid_target, draft_cnt)) temp = 0 for i in range(max_frame - 1, -1, -1): if appear[i, max_instance] != 0: temp = i break predict(max_frame - 1, temp, -1, [max_instance]) temp = frames_num for i in range(max_frame + 1, frames_num, 1): if appear[i, max_instance] != 0: temp = i break predict(max_frame + 1, temp, 1, [max_instance]) update_appear() for th in range(frames_num): save_frame(th, False, '%s_%05d' % (reid_target, draft_cnt)) draft_cnt = draft_cnt + 1 for th in range(frames_num): save_frame(th, False, '%s' % reid_target) if use_cache: pickle_dump((pred_prob, bbox_cnt), cache_file) for th in range(frames_num): save_frame(th, False, 'result', False)
def train(): global args, cfg, device args = parse_args() cfg = Config.from_file(args.config) #====================================================================== # ### Set up training # #====================================================================== # Set ?PU device cuda = torch.cuda.is_available() if cuda: print('\033[1m\033[91m' + '# cuda available!' + '\033[0m') device = torch.device(f'cuda:{args.gpu}') else: device = 'cpu' # set start iteration iteration = 0 # Set up networks to train num_class = len(cfg.train.dataset.class_list) gen = getattr(models, cfg.models.generator.model)(cfg.models.generator, num_class).to(device) dis = getattr(models, cfg.models.discriminator.model)(cfg.models.discriminator, cfg.train.dataset.frame_nums//cfg.train.dataset.frame_step, num_class).to(device) networks = {'gen': gen, 'dis': dis} # Load resume state_dict (to restart training) if args.resume: checkpoint_path = args.resume if os.path.isfile(checkpoint_path): print(f'loading checkpoint from {checkpoint_path}') checkpoint = torch.load(checkpoint_path, map_location=device) for name, model in networks.items(): model.load_state_dict(checkpoint[f'{name}_state_dict']) iteration = checkpoint['iteration'] # Set up an optimizer gen_lr = cfg.train.parameters.g_lr dis_lr = cfg.train.parameters.d_lr opts = {} opts['gen'] = torch.optim.Adam(gen.parameters(), lr=gen_lr, betas=(0.5, 0.999)) opts['dis'] = torch.optim.Adam(dis.parameters(), lr=dis_lr, betas=(0.5, 0.999)) # Load resume state_dict if args.resume: opts['gen'].load_state_dict(checkpoint['opt_gen_state_dict']) opts['dis'].load_state_dict(checkpoint['opt_dis_state_dict']) # Set up dataset train_dataset = BVHDataset(cfg.train.dataset, mode='train') train_loader = torch.utils.data.DataLoader( train_dataset, batch_size = cfg.train.batchsize, num_workers = cfg.train.num_workers, shuffle=True, drop_last=True) print(f'Data root \033[1m\"{cfg.train.dataset.data_root}\"\033[0m contains \033[1m{len(train_dataset)}\033[0m samples.') # Save scripts and command if not os.path.exists(cfg.train.out): os.makedirs(cfg.train.out) shutil.copy(args.config, f'./{cfg.train.out}') shutil.copy('./core/models/MotionGAN.py', f'./{cfg.train.out}') shutil.copy('./train.py', f'./{cfg.train.out}') commands = sys.argv with open(f'./{cfg.train.out}/command.txt', 'w') as f: f.write(f'python {commands[0]} ') for command in commands[1:]: f.write(command + ' ') # Set Criterion if cfg.train.GAN_type == 'normal': GAN_criterion = torch.nn.BCELoss().to(device) elif cfg.train.GAN_type == 'ls': GAN_criterion = torch.nn.MSELoss().to(device) else: GAN_criterion = None BCE_criterion = torch.nn.BCELoss().to(device) base_criterion = torch.nn.MSELoss().to(device) # Tensorboard Summary Writer writer = tbx.SummaryWriter(log_dir=os.path.join(cfg.train.out, 'log')) # train print('\033[1m\033[93m## Start Training!! ###\033[0m') while iteration < cfg.train.total_iterations: iteration = train_loop(train_loader, train_dataset, networks, opts, iteration, cfg.train.total_iterations, GAN_criterion, BCE_criterion, base_criterion, writer) # Save final model state = {'iteration':iteration, 'config':dict(cfg)} state[f'gen_state_dict'] = gen.state_dict() state[f'dis_state_dict'] = dis.state_dict() state['opt_gen_state_dict'] = opts['gen'].state_dict() state['opt_dis_state_dict'] = opts['dis'].state_dict() path = os.path.join(os.path.join(cfg.train.out,'checkpoint'), f'checkpoint.pth.tar') torch.save(state, path) torch.save(gen.state_dict(), os.path.join(cfg.train.out,f'gen.pth')) torch.save(dis.state_dict(), os.path.join(cfg.train.out,f'dis.pth')) print(f'trained model saved!') writer.close()
def main(): args = parse_args() cfg = Config.from_file(args.config) # Set ?PU device cuda = torch.cuda.is_available() if cuda: print('\033[1m\033[91m' + '# cuda available!' + '\033[0m') device = torch.device(f'cuda:{args.gpu}') else: device = 'cpu' ## Set up generator network num_class = len(cfg.train.dataset.class_list) gen = getattr(models, cfg.models.generator.model)(cfg.models.generator, num_class).to(device) ## Load weight if args.weight is not None: checkpoint_path = args.weight else: checkpoint_path = os.path.join(cfg.test.out, 'gen.pth') if not os.path.exists(checkpoint_path): checkpoint_path = sorted( glob.glob( os.path.join(cfg.test.out, 'checkpoint', 'iter_*.pth.tar')))[-1] if not os.path.exists(checkpoint_path): print('Generator weight not found!') sys.exit() else: print(f'Loading generator model from \033[1m{checkpoint_path}\033[0m') checkpoint = torch.load(checkpoint_path, map_location=device) if 'gen_state_dict' in checkpoint: gen.load_state_dict(checkpoint['gen_state_dict']) iteration = checkpoint['iteration'] else: gen.load_state_dict(checkpoint) iteration = cfg.train.total_iterations gen.eval() ## Create output directory result_dir = f'{cfg.test.out}/statistics/iter_{iteration}' if not os.path.exists(result_dir): os.makedirs(result_dir) ## Parse component list (for only PCA and tSNE) components = [] for pattern in re.findall(r'(\d+-\d+)', args.components): pair = pattern.split('-') components.append([int(pair[0]), int(pair[1])]) ## Conduct analisis # Apply PCA if args.mode == 'pca': result_path = os.path.join(result_dir, f'PCA_{args.target}.png') apply_pca(cfg, gen, result_path, args.target, components, device) # Apply tSNE elif args.mode == 'tSNE': result_path = os.path.join(result_dir, f'tSNE_{args.target}.pdf') apply_tSNE(cfg, gen, result_path, args.target, components, device) # Calcurate distance between centroid of each cluster and visualize as heatmap on matrix elif args.mode == 'heatmap': assert args.target == 'w' draw_heatmap(cfg, gen, result_dir, device) else: raise ValueError('Invalid mode!')
def make_video(): args = parse_args() cfg = Config.from_file(args.config) #====================================================================== # ### Set up model # #====================================================================== ## Set ?PU device cuda = torch.cuda.is_available() if cuda and args.gpu > -1: print('\033[1m' + '# cuda available!' + '\033[0m') device = torch.device(f'cuda:{args.gpu}') else: device = 'cpu' ## Define Generator gen = getattr(models, cfg.models.generator.model)(cfg.models.generator, len(cfg.train.dataset.class_list)).to(device) # Load weight if args.weight is None: checkpoint_path = os.path.join(cfg.test.out, 'gen.pth') if not os.path.exists(checkpoint_path): checkpoint_path = sorted(glob.glob(os.path.join(cfg.test.out, 'checkpoint', 'iter_*.pth.tar')))[-1] else: checkpoint_path = args.weight if not os.path.exists(checkpoint_path): print('\033[31m' + 'generator weight not found!' + '\033[0m') else: print('\033[33m' + 'loading generator model from ' + '\033[1m' + checkpoint_path + '\033[0m') checkpoint = torch.load(checkpoint_path, map_location=device) if 'gen_state_dict' in checkpoint: gen.load_state_dict(checkpoint['gen_state_dict']) iteration = checkpoint['iteration'] else: gen.load_state_dict(checkpoint) iteration = cfg.train.total_iterations gen.eval() #====================================================================== # ### Label Embedding # #====================================================================== ''' * Style Mode - single : Inference with each style - random : Randomly chosen 4 style - mix : Interpolation between two style - operation : Addition and subtraction with multiple style ''' ## Define label embed function if args.model in ['MotionGAN', 'TCN']: z = 0 * gen.make_hidden(1, 1).to(device) if cfg.models.generator.use_z else None LabelEmbedder = lambda style: gen.latent_transform(z, style_to_label(style)) elif args.model == 'ACGAN': z = 0 * gen.make_hidden(1, 1).to(device) if cfg.train.generator.use_z else None LabelEmbedder = lambda style: gen.label_emb(style_to_label(style)) elif args.model == 'LSTM': LabelEmbedder = lambda style: gen.label_emb(style_to_label(style)) else: pass # Define style name convert function def style_to_label(style): return torch.Tensor([cfg.train.dataset.class_list.index(style)]).type(torch.LongTensor).to(device) ## Create name and embed w pair inputs = [] # Mode : Single style inference if args.style_mode == 'single': style_list = args.style.split(',') if args.style else cfg.train.dataset.class_list for style in style_list: w = LabelEmbedder(style) inputs.append((style, w)) # Mode : Random style inference if args.style_mode == 'random': style_list = cfg.train.dataset.class_list for style in random.sample(style_list, args.num_samples): w = LabelEmbedder(style) inputs.append((style, w)) # Mode : Style Mix elif args.style_mode == 'mix': style1, style2 = args.style.split('-') w1 = LabelEmbedder(style1) w2 = LabelEmbedder(style2) inputs.append((style1, w1)) for ratio in [0.5]: w = ratio * w1 + (1.0 - ratio) * w2 inputs.append(('Mix'+str(ratio), w)) inputs.append((style2, w2)) # Mode : Style Operation elif args.style_mode == 'operation': def operation_parser(op_str): operands, operators = [], ['+'] tail = 0 for i in range(len(op_str)): if op_str[i] in ['+', '-']: operands.append(op_str[tail:i]) operators.append(op_str[i]) tail = i+1 operands.append(op_str[tail:]) assert len(operands) == len(operators) return operands, operators operands, operators = operation_parser(args.style) # Embed first operand w_result = 0 # Embed rest operands and calcurate operation for operand, operator in zip(operands, operators): w = LabelEmbedder(operand) inputs.append((operand,w)) if operator == '+': w_result += w elif operator == '-': w_result -= w else: raise ValueError('Invalid operator {operator}') inputs.append((args.style, w_result)) #====================================================================== # ### Define target data # #====================================================================== ''' - dataset : Inference on test dataset specified in config - draw : Inference on draw curve (.pkl file) ''' targets = [] ### Mode : Inference on test dataset if args.target == 'dataset': print(f'\033[92mInference on \033[1m{cfg.test.dataset.data_root}\033[0m') # Set up dataset test_dataset = BVHDataset(cfg.test.dataset, mode='test') # Prepare target data for k in range(len(test_dataset)): x_data, control_data, _ = test_dataset[k] x_data = torch.from_numpy(x_data) control_data = torch.from_numpy(control_data) # Convert tensor to batch x_data = x_data.unsqueeze(0).unsqueeze(1).type(torch.FloatTensor) control_data = control_data.unsqueeze(0).unsqueeze(1).type(torch.FloatTensor) original_curve = control_data[0,0,:,:].data.numpy() if not args.hide_curve else None # Generate input velocity spline v_control = control_data[:,:,1:,] - control_data[:,:,:-1,:] v_control = F.pad(v_control, (0,0,1,0), mode='reflect') v_control = v_control.to(device) targets.append({'name':f'{k:03d}', 'x_data':x_data, 'v_control':v_control, 'original_curve':original_curve}) ### Mode : Inference on draw curve elif args.target == 'draw': print(f'\033[92mInference on \033[1m{args.target_file}\33[0m') # Open .pkl file assert (args.target_file).endswith('.pkl') with open(args.target_file, mode='rb') as f: data = pickle.load(f) if 'trajectory' in data: trajectory_raw = data['trajectory'] control_point_interval = len(trajectory_raw) // args.control_point if args.control_point is not None else 120 control_data = convert_event_to_spline(trajectory_raw, control_point_interval=control_point_interval) data['control'] = control_data else: control_data = data['spline'] trajectory_raw = [] data_name = os.path.splitext(os.path.split(args.target_file)[1])[0] speed_modify = args.speed frame_step = int(cfg.train.frame_step // args.splinex * args.fps) control_data = torch.from_numpy(control_data).unsqueeze(0).unsqueeze(1).type(torch.FloatTensor) speed_modify /= scale batch_length = int(control_data.shape[2]/(frame_step*args.fps)) // 64 * 64 control_data = F.interpolate(control_data, size=(batch_length, spline_data.shape[3]), mode='bilinear', align_corners=True) control_data *= args.splinex / speed_modify if args.hide_curve: original_curve = None else: original_curve = (control_data[0,0,:,:] - control_data[0,0,0,:]).data.numpy() original_curve[:,1] = original_curve[:,1] - original_curve[:,1] # Convert position to velocity v_control = control_data[:,:,1:,] - control_data[:,:,:-1,:] v_control = F.pad(v_control, (0,0,1,0), mode='reflect') v_control = Variable(v_control).to(device) targets.append({'name': data_name, 'v_control':v_control, 'original_curve':original_curve}) #====================================================================== # ### Test Start # #====================================================================== ## Define output directory if args.target == 'dataset': test_dataset_name = os.path.split(cfg.test.dataset.data_root)[1] elif args.target == 'draw': test_dataset_name = 'Draw' result_dir_top = f'{cfg.test.out}/test/iter_{iteration}/{test_dataset_name}' if args.out is None else os.path.join(args.out, test_dataset_name) if not os.path.exists(result_dir_top): os.makedirs(result_dir_top) ## Testing option standard_bvh = cfg.test.dataset.standard_bvh if hasattr(cfg.test.dataset, 'standard_bvh') else 'core/datasets/CMU_standard.bvh' # Prepare initial pose for lstm if args.model == 'LSTM': if 'initial_pose' in checkpoint: initial_pose = checkpoint['initial_pose'] else: initial_pose = np.load('core/utils/initial_post_lstm.npy') initial_pose = torch.from_numpy(initial_pose).view(1,1,1,-1).type(torch.FloatTensor) ## Generate each sample for test_data in targets: v_control = test_data['v_control'] original_curve = test_data['original_curve'] result_list = [] for name, w in inputs: start_time = time.time() original_curve_j = original_curve.copy() #---------------------------------------- # Inference with model #---------------------------------------- ## MotionGAN if args.model == 'MotionGAN': fake_v_trajectory, x_fake_motion = gen(v_control, w=w) ## ACGAN elif args.model == 'ACGAN': fake_v_trajectory, x_fake_motion = gen(v_control, z.repeat(1,1,v_spline.shape[2],1), label_embed=w) ## TCN elif args.model == 'TCN': # Inference each 128 frames for t in range(0, v_spline.shape[2], 128): fake_v_trajectory_t, x_fake_motion_t = gen(v_control[:,:,t:t+128,:], w=w) fake_v_trajectory = fake_v_trajectory_t if t==0 else torch.cat((fake_v_trajectory, fake_v_trajectory_t), dim=2) x_fake_motion = x_fake_motion_t if t==0 else torch.cat((x_fake_motion, x_fake_motion_t), dim=2) ## LSTM elif args.model == 'LSTM': traj_t, pose_t = v_control[:,:,:1,:], initial_pose.to(device) for t in range(v_control.shape[2]): traj_t, pose_t = gen(v_control[:,:,t,:], pose_t, traj_t, label_embed=w) fake_v_trajectory = traj_t if t==0 else torch.cat((fake_v_trajectory, traj_t), dim=2) x_fake_motion = pose_t if t==0 else torch.cat((x_fake_motion, pose_t), dim=2) #--------------------------------------------------- # Convert model output to viewable joint position #--------------------------------------------------- if x_fake_motion.shape[2] > args.start_frame: x_fake_motion = x_fake_motion[:,:,args.start_frame:,:] fake_v_trajectory = fake_v_trajectory[:,:,args.start_frame:,:] if original_curve_j is not None: original_curve_j = original_curve_j[args.start_frame:,:] - original_curve_j[args.start_frame,:] else: if original_curve_j is not None: original_curve_j = original_curve_j - original_curve_j[0,:] # Root position at start frame start_position = torch.zeros(1,1,1,3) if re.search(r"OldMan|Sneaky|Scared|Chicken|Dinosaur", name) is not None: start_position[0,0,0,1] = 15.0 / cfg.test.dataset.scale else: start_position[0,0,0,1] = 17.0 / cfg.test.dataset.scale # Velocity to positon fake_trajectory = reconstruct_v_trajectory(fake_v_trajectory.data.cpu(), start_position) x_fake = torch.cat((fake_trajectory, x_fake_motion.cpu()), dim=3) result_list.append({'caption': name, 'motion': x_fake.detach()[0,:,:,:], 'control': original_curve_j}) # Measure time avg_time = (time.time() - start_time) #------------------------------------------------ # Save each sample #------------------------------------------------ if args.save_separate: if args.style_mode == 'single': result_dir = os.path.join(result_dir_top, name) else: result_dir = os.path.join(result_dir_top, args.style) if not os.path.exists(result_dir): os.mkdir(result_dir) result_path = os.path.join(result_dir, test_data['name']+'_'+name) print(f'\nInference : {result_path} {x_fake[0,0,:,:].shape} Time: {avg_time:.05f}') # Save result data with open(result_path+'.pkl', 'wb') as f: pickle.dump(result_dic, f) # Save video save_video(result_path+'.'+args.save_format, result_list, cfg.test, camera_move='stand', elev=args.elev, azim=args.azim) if args.save_pics: save_timelapse(result_path+'.png', result_list, cfg.test) result_list = [] #------------------------------------------------ # Save all frame in one video #------------------------------------------------ if not args.save_separate: result_dir = os.path.join(result_dir_top, args.style) if not os.path.exists(result_dir): os.mkdir(result_dir) result_path = os.path.join(result_dir, test_data['name']) print(f'\nInference : {result_path} {x_fake[0,0,:,:].shape} Time: {avg_time:.05f}') # Save video save_video(result_path+'.'+args.save_format, result_list, cfg.test, camera_move='stand', elev=args.elev, azim=args.azim) if args.save_pics: save_timelapse(result_path+'.png', result_list, cfg.test)
def calcurate_style_meanstd(): parser = argparse.ArgumentParser() parser.add_argument('config') parser.add_argument('--out', required=True) args = parser.parse_args() cfg = Config.from_file(args.config) npy_paths = collect_path(cfg.train.dataset) head, ext = os.path.splitext(cfg.train.dataset) head, data_name = os.path.split(head) # Prepare skelton skelton, non_end_bones, joints_to_index, permute_xyz_order = btoj.get_standard_format( cfg.standard_bvh) _, non_zero_joint_to_index = btoj.cut_zero_length_bone( skelton, joints_to_index) Lul_offset = np.array(skelton['LeftUpLeg']['offsets']) Rul_offset = np.array(skelton['RightUpLeg']['offsets']) Lul_index = non_zero_joint_to_index['LeftUpLeg'] Rul_index = non_zero_joint_to_index['RightUpLeg'] standard_vector = Lul_offset - Rul_offset standard_norm = np.sqrt((Lul_offset[0] - Rul_offset[0])**2 + (Lul_offset[2] - Rul_offset[2])**2) # Initialize result = { style: {joint: 0 for joint in non_zero_joint_to_index} for style in cfg.train.class_list } data = None for npy_path in npy_paths: motion = np.load(npy_path) # Cut zero from motion _, motion = btoj.cut_zero_length_bone_frames(motion, skelton, joints_to_index) # Convert trajectory to velocity motion = motion[1:] trajectory = motion[:, :3] velocity = trajectory[1:, :] - trajectory[:-1, :] motion = np.concatenate((velocity, motion[1:, 3:]), axis=1) # Get orientation ('xz' only) motion_oriented = np.zeros_like(motion) leftupleg = motion[:, Lul_index * 3:Lul_index * 3 + 3] rightupleg = motion[:, Rul_index * 3:Rul_index * 3 + 3] vector = leftupleg - rightupleg norm = np.sqrt(vector[:, 0]**2 + vector[:, 2]**2) cos = (vector[:, 0] * standard_vector[0] + vector[:, 2] * standard_vector[2]) / (norm * standard_norm) cos = np.clip(cos, -1, 1) sin = 1 - cos**2 for t in range(motion.shape[0]): rotation_mat = np.array([[cos[t], 0., -sin[t]], [0., 1., 0.], [sin[t], 0., cos[t]]]) motion_oriented[t, :] = np.dot( rotation_mat.T, motion[t, :].reshape(28, 3).T).T.reshape(-1, ) # Set class npy_name = os.path.splitext(os.path.split(npy_path)[1])[0] style = npy_name.split('_')[0] mean = np.mean(motion[1:], axis=0) std = np.std(motion[1:], axis=0) mean_oriented = np.mean(motion_oriented, axis=0) std_oriented = np.std(motion_oriented, axis=0) # Write for joint in non_zero_joint_to_index: ji = non_zero_joint_to_index[joint] result[style][joint] = { 'mean': mean_oriented[ji * 3:ji * 3 + 3], 'std': std_oriented[ji * 3:ji * 3 + 3] } with open(args.out, 'wb') as f: pickle.dump(result, f) ##t-SNE mean_data, std_data, label_data = [], [], [] for style in result.keys(): data = result[style] mean_data.append( [data[joint]['mean'] for joint in non_zero_joint_to_index]) std_data.append( [data[joint]['std'] for joint in non_zero_joint_to_index]) label_data.append(cfg.train.class_list.index(style)) mean_data = np.array(mean_data).reshape(len(mean_data), -1) std_data = np.array(std_data).reshape(len(mean_data), -1) label_data = np.array(label_data).reshape(len(mean_data), ) plt.figure(figsize=(30, 30), dpi=72) # joint mean mean_velocity = np.stack((mean_data[:, 6], mean_data[:, 7]), axis=1) plt.subplot(331) plt.scatter(mean_velocity[:, 0], mean_velocity[:, 1], s=25) plt.title('mean' + list(non_zero_joint_to_index.keys())[2]) for i in range(mean_velocity.shape[0]): point = mean_velocity[i] label = cfg.train.class_list[label_data[i]] plt.text(point[0], point[1], label, fontsize=8) # joint std std_velocity = np.stack((std_data[:, 6], std_data[:, 7]), axis=1) plt.subplot(332) plt.scatter(std_velocity[:, 0], std_velocity[:, 1], s=25) plt.title('std_' + list(non_zero_joint_to_index.keys())[2]) for i in range(std_velocity.shape[0]): point = std_velocity[i] label = cfg.train.class_list[label_data[i]] plt.text(point[0], point[1], label, fontsize=8) # PCA mean pca = PCA(n_components=2) mean_pca = pca.fit_transform(mean_data) plt.subplot(333) plt.scatter(mean_pca[:, 0], mean_pca[:, 1], s=25) plt.title('pca_mean') for i in range(mean_pca.shape[0]): point = mean_pca[i] label = cfg.train.class_list[label_data[i]] plt.text(point[0], point[1], label, fontsize=8) # joint mean mean_velocity = np.stack((mean_data[:, 36], mean_data[:, 37]), axis=1) plt.subplot(334) plt.scatter(mean_velocity[:, 0], mean_velocity[:, 1], s=25) plt.title('mean_' + list(non_zero_joint_to_index.keys())[12]) for i in range(mean_velocity.shape[0]): point = mean_velocity[i] label = cfg.train.class_list[label_data[i]] plt.text(point[0], point[1], label, fontsize=8) # joint std std_velocity = np.stack((std_data[:, 36], std_data[:, 37]), axis=1) plt.subplot(335) plt.scatter(std_velocity[:, 0], std_velocity[:, 1], s=25) plt.title('std_' + list(non_zero_joint_to_index.keys())[12]) for i in range(std_velocity.shape[0]): point = std_velocity[i] label = cfg.train.class_list[label_data[i]] plt.text(point[0], point[1], label, fontsize=8) # PCA mean pca = PCA(n_components=2) std_pca = pca.fit_transform(std_data) plt.subplot(336) plt.scatter(std_pca[:, 0], std_pca[:, 1], s=25) plt.title('pca_std') for i in range(std_pca.shape[0]): point = std_pca[i] label = cfg.train.class_list[label_data[i]] plt.text(point[0], point[1], label, fontsize=8) # joint mean mean_velocity = np.stack((mean_data[:, 60], mean_data[:, 61]), axis=1) plt.subplot(337) plt.scatter(mean_velocity[:, 0], mean_velocity[:, 1], s=25) plt.title('mean_' + list(non_zero_joint_to_index.keys())[20]) for i in range(mean_velocity.shape[0]): point = mean_velocity[i] label = cfg.train.class_list[label_data[i]] plt.text(point[0], point[1], label, fontsize=8) # joint std std_velocity = np.stack((std_data[:, 60], std_data[:, 61]), axis=1) plt.subplot(338) plt.scatter(std_velocity[:, 0], std_velocity[:, 1], s=25) plt.title('std_' + list(non_zero_joint_to_index.keys())[20]) for i in range(std_velocity.shape[0]): point = std_velocity[i] label = cfg.train.class_list[label_data[i]] plt.text(point[0], point[1], label, fontsize=8) # PCA mean and std pca = PCA(n_components=2) mean_std_pca = pca.fit_transform( np.concatenate((mean_data, std_data), axis=1)) plt.subplot(339) plt.scatter(mean_std_pca[:, 0], mean_std_pca[:, 1], s=25) plt.title('pca_mean_std') for i in range(mean_std_pca.shape[0]): point = mean_std_pca[i] label = cfg.train.class_list[label_data[i]] plt.text(point[0], point[1], label, fontsize=8) plt.savefig(os.path.splitext(args.out)[0] + '_tSNE.png') # 3D PCA mean and std fig = plt.figure(figsize=(10, 10), dpi=72) ax = Axes3D(fig) pca = PCA(n_components=3) mean_std_pca3d = pca.fit_transform( np.concatenate((mean_data, std_data), axis=1)) ax.scatter3D(mean_std_pca3d[:, 0], mean_std_pca3d[:, 1], mean_std_pca3d[:, 2], s=25) for i in range(mean_std_pca.shape[0]): point = mean_std_pca3d[i] label = cfg.train.class_list[label_data[i]] ax.text(point[0], point[1], point[2], label, fontsize=8) plt.savefig(os.path.splitext(args.out)[0] + '_PCA3D.png')
def main(): args = parse_args() cfg = Config.from_file(args.config) ## Set ?PU device cuda = torch.cuda.is_available() if cuda and args.gpu > -1: print('\033[1m\033[91m' + '# cuda available!' + '\033[0m') device = torch.device(f'cuda:{args.gpu}') else: device = 'cpu' # Set up generator network num_class = len(cfg.train.dataset.class_list) gen = getattr(models, cfg.models.generator.model)(cfg.models.generator, num_class).to(device) # Load weight if args.weight is None: checkpoint_path = os.path.join(cfg.test.out, 'gen.pth') if not os.path.exists(checkpoint_path): checkpoint_path = sorted( glob.glob( os.path.join(cfg.test.out, 'checkpoint', 'iter_*.pth.tar')))[-1] else: checkpoint_path = args.weight if not os.path.exists(checkpoint_path): print('Generator weight not found!') else: print(f'Loading generator model from \033[1m{checkpoint_path}\033[0m') checkpoint = torch.load(checkpoint_path, map_location=device) if 'gen_state_dict' in checkpoint: gen.load_state_dict(checkpoint['gen_state_dict']) iteration = checkpoint['iteration'] else: gen.load_state_dict(checkpoint) iteration = cfg.train.total_iterations gen.eval() ## Create name and embed w pair inputs = [] z = gen.make_hidden(1, 1).to(device) if cfg.models.generator.use_z else None for i, style in enumerate(cfg.train.dataset.class_list): label = torch.Tensor([i]).type(torch.LongTensor).to(device) w = gen.latent_transform(z, label) inputs.append((style, w)) # Each label corresponds to rows rows = ['Average'] + [s[0] for s in inputs] #====================================================================== # ### Prepare target data # #====================================================================== print(f'Inference on \033[1m{cfg.test.dataset.data_root}\033[0m') targets = [] # Set up dataset test_dataset = BVHDataset(cfg.test.dataset, mode='test') # Prepare target data for k in range(len(test_dataset)): x_data, control_data, _ = test_dataset[k] x_data = torch.from_numpy(x_data) control_data = torch.from_numpy(control_data) # Convert tensor to batch x_data = x_data.unsqueeze(0).unsqueeze(1).type(torch.FloatTensor) control_data = control_data.unsqueeze(0).unsqueeze(1).type( torch.FloatTensor) original_curve = control_data[0, 0, :, :] # Generate input velocity spline v_control = control_data[:, :, 1:, ] - control_data[:, :, :-1, :] v_control = F.pad(v_control, (0, 0, 1, 0), mode='reflect') v_control = v_control.to(device) targets.append({ 'name': f'{k:03d}', 'x_data': x_data, 'v_control': v_control, 'original_curve': original_curve }) # Each target data corresponds to columns columns = ['Average'] + [data['name'] for data in targets] #====================================================================== # ### Test Start # #====================================================================== ## Define output directory test_dataset_name = os.path.split(cfg.test.dataset.data_root)[1] result_dir = f'{cfg.test.out}/eval/iter_{iteration}/{test_dataset_name}' if not os.path.exists(result_dir): os.makedirs(result_dir) ## Testing option standard_bvh = cfg.test.dataset.standard_bvh if hasattr( cfg.test.dataset, 'standard_bvh') else 'core/datasets/CMU_standard.bvh' skelton, _, joints_to_index, _ = get_standard_format(standard_bvh) _, non_zero_joint_to_index = cut_zero_length_bone(skelton, joints_to_index) id_rightleg = non_zero_joint_to_index['RightToeBase'] id_leftleg = non_zero_joint_to_index['LeftToeBase'] ## Evaluate each sample trajectory_error_data = np.zeros((len(rows), len(columns))) footskate_data = np.zeros((len(rows), len(columns))) for i, test_data in enumerate(targets): v_control = test_data['v_control'] original_curve = test_data['original_curve'] original_curve = original_curve - original_curve[0, :] result_dic = {} for j, (name, w) in enumerate(inputs): start_time = time.time() #---------------------------------------- # Inference with model #---------------------------------------- fake_v_trajectory, x_fake = gen(v_control, w=w) # Velocity to positon fake_trajectory = reconstruct_v_trajectory( fake_v_trajectory.data.cpu(), torch.zeros(1, 1, 1, 3)) x_fake = torch.cat((fake_trajectory, x_fake.cpu()), dim=3) # Denormalize x_fake *= cfg.test.dataset.scale #--------------------------------------------------- # Calcurate on metrics #--------------------------------------------------- frame_length = x_fake.shape[2] # Calcurlate foot skating distance footskate_dist = calcurate_footskate( x_fake[0, 0, :, :3], x_fake[0, 0, :, id_rightleg * 3:id_rightleg * 3 + 3], x_fake[0, 0, :, id_leftleg * 3:id_leftleg * 3 + 3]) / frame_length footskate_data[j + 1, i + 1] = round(footskate_dist, 6) # Calcurlate trajectory error error_dist = calcurate_trajectory_error(x_fake[0, 0, :, :3], original_curve, 8, 32) trajectory_error_data[j + 1, i + 1] = round(error_dist, 6) #--------------------------------------------------- # Merge all results #--------------------------------------------------- for name, data in [('trajectory_erorr_dist', trajectory_error_data), ('footskate_dist', footskate_data)]: # Get average data[:, 0] = np.sum(data[:, 1:], axis=1) / (len(columns) - 1) data[0, :] = np.sum(data[1:, :], axis=0) / (len(rows) - 1) data = data.tolist() # Save as csv df = pd.DataFrame(data, index=rows, columns=columns) print(name, '\n', df) df.to_csv(os.path.join(result_dir, f'{name}.csv'))
def test(): global args, cfg, device args = parse_args() cfg = Config.from_file(args.config) # Set ?PU device cuda = torch.cuda.is_available() if cuda: print('\033[1m\033[91m' + '# cuda available!' + '\033[0m') device = torch.device(f'cuda:{args.gpu}') else: device = 'cpu' ##################################################### ## Prepare for test ##################################################### # Set up generator network num_class = len(cfg.train.dataset.class_list) gen = getattr(models, cfg.models.generator.model)(cfg.models.generator, num_class).to(device) total_params = sum(p.numel() for p in gen.parameters() if p.requires_grad) print(f'Total parameter amount : \033[1m{total_params}\033[0m') # Load weight if args.weight is not None: checkpoint_path = args.weight else: checkpoint_path = os.path.join(cfg.test.out, 'gen.pth') if not os.path.exists(checkpoint_path): checkpoint_path = sorted( glob.glob( os.path.join(cfg.test.out, 'checkpoint', 'iter_*.pth.tar')))[-1] if not os.path.exists(checkpoint_path): print('Generator weight not found!') else: print(f'Loading generator model from \033[1m{checkpoint_path}\033[0m') checkpoint = torch.load(checkpoint_path, map_location=device) if 'gen_state_dict' in checkpoint: gen.load_state_dict(checkpoint['gen_state_dict']) iteration = checkpoint['iteration'] else: gen.load_state_dict(checkpoint) iteration = cfg.train.total_iterations gen.eval() # Set up dataset test_dataset = BVHDataset(cfg.test.dataset, mode='test') test_dataset_name = os.path.split( cfg.test.dataset.data_root.replace('*', ''))[1] # Set standard bvh standard_bvh = cfg.test.dataset.standard_bvh if hasattr( cfg.test.dataset, 'standard_bvh') else 'core/datasets/CMU_standard.bvh' # Create output directory result_dir = f'{cfg.test.out}/test/iter_{iteration}/{test_dataset_name}' if not os.path.exists(result_dir): os.makedirs(result_dir) ##################################################### ## Test start ##################################################### for i in range(len(test_dataset)): x_data, control_data, label = test_dataset[i] if x_data.shape[ 0] < cfg.train.dataset.frame_nums // cfg.train.dataset.frame_step: continue # Motion and control signal data x_data = torch.from_numpy(x_data).unsqueeze(0).unsqueeze(1).type( torch.FloatTensor) x_real = Variable(x_data).to(device) control_data = torch.from_numpy(control_data).unsqueeze(0).unsqueeze( 1).type(torch.FloatTensor) control = control_data.to(device) # Convert root trajectory to verocity gt_trajectory = x_data[:, :, :, 0:3] gt_v_trajectory = gt_trajectory[:, :, 1:, :] - gt_trajectory[:, :, :-1, :] gt_v_trajectory = F.pad(gt_v_trajectory, (0, 0, 1, 0), mode='reflect') gt_v_trajectory = Variable(gt_v_trajectory).to(device) # Convert control curve to velocity v_control = control[:, :, 1:, ] - control[:, :, :-1, :] v_control = F.pad(v_control, (0, 0, 1, 0), mode='reflect') v_control = Variable(v_control).to(device) results_list = [] start_time = time.time() # Generate fake sample for k in range(args.num_samples): # Generate noize z z = gen.make_hidden(1, x_data.shape[2]).to( device) if cfg.models.generator.use_z else None fake_label = torch.randint(0, len(cfg.train.dataset.class_list), size=(1, )).type( torch.LongTensor).to(device) fake_v_trajectory, x_fake = gen(v_control, z, fake_label) fake_trajectory = reconstruct_v_trajectory( fake_v_trajectory.data.cpu(), torch.zeros(1, 1, 1, 3)) caption = f'{cfg.train.dataset.class_list[fake_label]}_{k}' results_list.append({ 'caption': caption, 'motion': torch.cat((fake_trajectory, x_fake.data.cpu()), dim=3), 'control': control.data.cpu() }) avg_time = (time.time() - start_time) / args.num_samples # Save results result_path = result_dir + f'/{i:03d}.avi' print( f'\nInference : {result_path} ({v_control.shape[2]} frames) Time: {avg_time:.05f}' ) save_video(result_path, results_list, cfg.test)
def main(): global cfg, best_PSNR args = parse_args() cfg = Config.from_file(args.config) os.environ["CUDA_VISIBLE_DEVICES"] = ','.join( str(gpu) for gpu in cfg.device) # os.environ["CUDA_VISIBLE_DEVICES"] = '0' cudnn.benchmark = True cudnn.fastest = True if hasattr(datasets, cfg.dataset): ds = getattr(datasets, cfg.dataset) else: raise ValueError('Unknown dataset ' + cfg.dataset) model = getattr(models, cfg.model.name)(cfg.model).cuda() cfg.train.input_mean = model.input_mean cfg.train.input_std = model.input_std cfg.test.input_mean = model.input_mean cfg.test.input_std = model.input_std # Data loading code # train_loader = torch.utils.data.DataLoader( # ds(cfg.train), # batch_size=cfg.train.batch_size, # shuffle=True, # num_workers=0, #32, # pin_memory=True, # drop_last=True) val_loader = torch.utils.data.DataLoader( datasets.UCF101(cfg.test, False), batch_size=cfg.test.batch_size, shuffle=False, num_workers=0, #32, pin_memory=True) cfg.train.optimizer.args.max_iter = ( cfg.train.optimizer.args.max_epoch * len(train_loader)) policies = model.get_optim_policies() for group in policies: print(('group: {} has {} params, lr_mult: {}, decay_mult: {}'.format( group['name'], len(group['params']), group['lr_mult'], group['decay_mult']))) optimizer = Optim(policies, cfg.train.optimizer) if cfg.resume or cfg.weight: checkpoint_path = cfg.resume if cfg.resume else cfg.weight if os.path.isfile(checkpoint_path): print(("=> loading checkpoint '{}'".format(checkpoint_path))) checkpoint = torch.load(checkpoint_path) model.load_state_dict(checkpoint['state_dict'], False) if cfg.resume: optimizer.load_state_dict(checkpoint['grad_dict']) else: print(("=> no checkpoint found at '{}'".format(checkpoint_path))) model = DataParallelwithSyncBN( model, device_ids=range(len(cfg.device))).cuda() # define loss function (criterion) optimizer and evaluator criterion = torch.nn.MSELoss().cuda() evaluator = EvalPSNR(255.0 / np.mean(cfg.test.input_std)) # PSNR = validate(val_loader, model, optimizer, criterion, evaluator) # return for epoch in range(cfg.train.optimizer.args.max_epoch): # train for one epoch # train(train_loader, model, optimizer, criterion, epoch) # evaluate on validation set if ((epoch + 1) % cfg.logging.eval_freq == 0 or epoch == cfg.train.optimizer.args.max_epoch - 1): PSNR = validate(val_loader, model, optimizer, criterion, evaluator) # remember best PSNR and save checkpoint is_best = PSNR > best_PSNR best_PSNR = max(PSNR, best_PSNR) save_checkpoint({ 'epoch': epoch + 1, 'arch': dict(cfg), 'state_dict': model.module.state_dict(), 'grad_dict': optimizer.state_dict(), 'best_PSNR': best_PSNR, }, is_best)