diff --git a/configs/pose3d/MB_ft_h36m_global_lite.yaml b/configs/pose3d/MB_ft_h36m_global_lite.yaml index 4a2d5d9..d42edb6 100644 --- a/configs/pose3d/MB_ft_h36m_global_lite.yaml +++ b/configs/pose3d/MB_ft_h36m_global_lite.yaml @@ -44,7 +44,7 @@ lambda_av: 0.0 # Augmentation synthetic: False -flip: True +flip: False mask_ratio: 0. mask_T_ratio: 0. noise: False \ No newline at end of file diff --git a/configs/pretrain/MB_lite.yaml b/configs/pretrain/MB_lite.yaml index ae06971..34250ab 100644 --- a/configs/pretrain/MB_lite.yaml +++ b/configs/pretrain/MB_lite.yaml @@ -1,17 +1,18 @@ # General -train_2d: True -no_eval: False -finetune: False +train_2d: False +no_eval: True +eval_last: False +finetune: True partial_train: null # Traning -epochs: 90 -checkpoint_frequency: 30 -batch_size: 64 +epochs: 10 +checkpoint_frequency: 1 +batch_size: 16 dropout: 0.0 -learning_rate: 0.0005 -weight_decay: 0.01 -lr_decay: 0.99 +learning_rate: 0.00025 +weight_decay: 0.02 +lr_decay: 0.98 pretrain_3d_curriculum: 30 # Model @@ -24,16 +25,16 @@ num_heads: 8 att_fuse: True # Data -data_root: data/motion3d/MB3D_f243s81/ -subset_list: [AMASS, H36M-SH] +data_root: ../DL_Project/pose_3d_v3/ +subset_list: [frame_81] dt_file: h36m_sh_conf_cam_source_final.pkl clip_len: 243 data_stride: 81 -rootrel: True +rootrel: False sample_stride: 1 num_joints: 17 no_conf: False -gt_2d: False +gt_2d: True # Loss lambda_3d_velocity: 20.0 @@ -44,10 +45,10 @@ lambda_a: 0.0 lambda_av: 0.0 # Augmentation -synthetic: True # synthetic: don't use 2D detection results, fake it (from 3D) -flip: True -mask_ratio: 0.05 -mask_T_ratio: 0.1 -noise: True +synthetic: False # synthetic: don't use 2D detection results, fake it (from 3D) +flip: False +mask_ratio: 0.0 +mask_T_ratio: 0 +noise: False noise_path: params/synthetic_noise.pth d2c_params_path: params/d2c_params.pkl \ No newline at end of file diff --git a/infer_wild.py b/infer_wild.py index 17acd19..baca6e4 100644 --- a/infer_wild.py +++ b/infer_wild.py @@ -12,6 +12,9 @@ from lib.data.dataset_wild import WildDetDataset from lib.utils.vismo import render_and_save +import os +os.environ['KMP_DUPLICATE_LIB_OK'] = 'TRUE' + def parse_args(): parser = argparse.ArgumentParser() parser.add_argument("--config", type=str, default="configs/pose3d/MB_ft_h36m_global_lite.yaml", help="Path to the config file.") @@ -28,23 +31,28 @@ def parse_args(): opts = parse_args() args = get_config(opts.config) +print('Loading checkpoint', opts.evaluate) +checkpoint = torch.load(opts.evaluate, map_location=lambda storage, loc: storage) + model_backbone = load_backbone(args) if torch.cuda.is_available(): model_backbone = nn.DataParallel(model_backbone) model_backbone = model_backbone.cuda() + model_backbone.load_state_dict(checkpoint['model_pos'], strict=True) +else: + state_dict = checkpoint['model_pos'] + state_dict = {k.replace('module.', '', 1): v for k, v in state_dict.items()} + model_backbone.load_state_dict(state_dict, strict=True) -print('Loading checkpoint', opts.evaluate) -checkpoint = torch.load(opts.evaluate, map_location=lambda storage, loc: storage) -model_backbone.load_state_dict(checkpoint['model_pos'], strict=True) model_pos = model_backbone model_pos.eval() testloader_params = { 'batch_size': 1, 'shuffle': False, - 'num_workers': 8, + 'num_workers': 0, 'pin_memory': True, - 'prefetch_factor': 4, - 'persistent_workers': True, + # 'prefetch_factor': 4, + # 'persistent_workers': True, 'drop_last': False } @@ -81,7 +89,8 @@ def parse_args(): if args.rootrel: predicted_3d_pos[:,:,0,:]=0 # [N,T,17,3] else: - predicted_3d_pos[:,0,0,2]=0 + # predicted_3d_pos[:,0,0,2]=0 + predicted_3d_pos[:, :, 0, 2] = 0 pass if args.gt_2d: predicted_3d_pos[...,:2] = batch_input[...,:2] diff --git a/infer_wild_mesh.py b/infer_wild_mesh.py index 0c1c9b7..7580b56 100644 --- a/infer_wild_mesh.py +++ b/infer_wild_mesh.py @@ -61,7 +61,10 @@ def solve_scale(x, y): # root_rel # args.rootrel = True -smpl = SMPL(args.data_root, batch_size=1).cuda() +if torch.cuda.is_available(): + smpl = SMPL(args.data_root, batch_size=1).cuda() +else: + smpl = SMPL(args.data_root, batch_size=1) J_regressor = smpl.J_regressor_h36m end = time.time() @@ -71,14 +74,20 @@ def solve_scale(x, y): model = MeshRegressor(args, backbone=model_backbone, dim_rep=args.dim_rep, hidden_dim=args.hidden_dim, dropout_ratio=args.dropout) print(f'init whole model time: {(time.time()-end):02f}s') +chk_filename = opts.evaluate if opts.evaluate else opts.resume +print('Loading checkpoint', chk_filename) +checkpoint = torch.load(chk_filename, map_location=lambda storage, loc: storage) + if torch.cuda.is_available(): model = nn.DataParallel(model) model = model.cuda() + model.load_state_dict(checkpoint['model'], strict=True) +else: + state_dict = checkpoint['model'] + state_dict = {k.replace('module.', '', 1): v for k, v in state_dict.items()} + model.load_state_dict(state_dict, strict=True) + -chk_filename = opts.evaluate if opts.evaluate else opts.resume -print('Loading checkpoint', chk_filename) -checkpoint = torch.load(chk_filename, map_location=lambda storage, loc: storage) -model.load_state_dict(checkpoint['model'], strict=True) model.eval() testloader_params = { diff --git a/lib/data/dataset_motion_3d.py b/lib/data/dataset_motion_3d.py index a2de10d..c8e8381 100644 --- a/lib/data/dataset_motion_3d.py +++ b/lib/data/dataset_motion_3d.py @@ -8,7 +8,7 @@ from torch.utils.data import Dataset, DataLoader from lib.data.augmentation import Augmenter3D from lib.utils.tools import read_pkl -from lib.utils.utils_data import flip_data +from lib.utils.utils_data import flip_data, crop_scale class MotionDataset(Dataset): def __init__(self, args, subset_list, data_split): # data_split: train/test @@ -38,13 +38,53 @@ def __init__(self, args, subset_list, data_split): self.synthetic = args.synthetic self.aug = Augmenter3D(args) self.gt_2d = args.gt_2d + + def __coco_to_h36m(self, kps_coco, n_dim=2): + """ + kps_coco: np.array of shape (N, 17, 2) in COCO format + returns: np.array of shape (N, 17, 2) in H36M format + """ + # this is the same function from `motionbert/motionbert_convert.ipynb` in the DL_project repo + N = kps_coco.shape[0] + h36m = np.zeros((N, 17, n_dim), dtype=np.float32) + + # Synthesize H36M joints not present in COCO + # H36M[0] = Hip root = average of left and right hips + h36m[:, 0] = (kps_coco[:, 11] + kps_coco[:, 12]) / 2 # (LHip + RHip) / 2 + + # H36M[7] = Spine mid = average of hip root and neck + neck = (kps_coco[:, 5] + kps_coco[:, 6]) / 2 # avg of shoulders as proxy + h36m[:, 7] = (h36m[:, 0] + neck) / 2 + + # H36M[8] = Thorax/Neck = average of shoulders + h36m[:, 8] = neck + + # H36M[10] = Head top = average of ears (or use nose as fallback) + h36m[:, 10] = (kps_coco[:, 3] + kps_coco[:, 4]) / 2 # avg of ears + + # Direct mappings (COCO index -> H36M index) + h36m[:, 1] = kps_coco[:, 12] # RHip + h36m[:, 2] = kps_coco[:, 14] # RKnee + h36m[:, 3] = kps_coco[:, 16] # RAnkle + h36m[:, 4] = kps_coco[:, 11] # LHip + h36m[:, 5] = kps_coco[:, 13] # LKnee + h36m[:, 6] = kps_coco[:, 15] # LAnkle + h36m[:, 9] = kps_coco[:, 0] # Nose + h36m[:, 11] = kps_coco[:, 5] # LShoulder + h36m[:, 12] = kps_coco[:, 7] # LElbow + h36m[:, 13] = kps_coco[:, 9] # LWrist + h36m[:, 14] = kps_coco[:, 6] # RShoulder + h36m[:, 15] = kps_coco[:, 8] # RElbow + h36m[:, 16] = kps_coco[:, 10] # RWrist + + return h36m def __getitem__(self, index): 'Generates one sample of data' # Select sample file_path = self.file_list[index] motion_file = read_pkl(file_path) - motion_3d = motion_file["data_label"] + motion_3d = self.__coco_to_h36m(motion_file["data_label"], 3) if self.data_split=="train": if self.synthetic or self.gt_2d: motion_3d = self.aug.augment3D(motion_3d) @@ -52,17 +92,19 @@ def __getitem__(self, index): motion_2d[:,:,:2] = motion_3d[:,:,:2] motion_2d[:,:,2] = 1 # No 2D detection, use GT xy and c=1. elif motion_file["data_input"] is not None: # Have 2D detection - motion_2d = motion_file["data_input"] + motion_2d = self.__coco_to_h36m(motion_file["data_input"], 3) + motion_2d = crop_scale(motion_2d) if self.flip and random.random() > 0.5: # Training augmentation - random flipping motion_2d = flip_data(motion_2d) motion_3d = flip_data(motion_3d) else: raise ValueError('Training illegal.') elif self.data_split=="test": - motion_2d = motion_file["data_input"] + motion_2d = self.__coco_to_h36m(motion_file["data_input"], 3) + motion_2d = crop_scale(motion_2d) if self.gt_2d: motion_2d[:,:,:2] = motion_3d[:,:,:2] motion_2d[:,:,2] = 1 else: - raise ValueError('Data split unknown.') + raise ValueError('Data split unknown.') return torch.FloatTensor(motion_2d), torch.FloatTensor(motion_3d) \ No newline at end of file diff --git a/lib/data/dataset_wild.py b/lib/data/dataset_wild.py index 8a462c4..3fd05d9 100644 --- a/lib/data/dataset_wild.py +++ b/lib/data/dataset_wild.py @@ -44,6 +44,9 @@ def halpe2h36m(x): {25, "RHeel"}, ''' T, V, C = x.shape + if V == 17: + return x + y = np.zeros([T,17,C]) y[:,0,:] = x[:,19,:] y[:,1,:] = x[:,12,:] diff --git a/lib/utils/vismo.py b/lib/utils/vismo.py index 456c3d7..4021ac4 100644 --- a/lib/utils/vismo.py +++ b/lib/utils/vismo.py @@ -263,12 +263,12 @@ def motion2video_3d(motion, save_path, fps=25, keep_imgs = False): ax.set_xlim(-512, 0) ax.set_ylim(-256, 256) ax.set_zlim(-512, 0) - # ax.set_xlabel('X') - # ax.set_ylabel('Y') - # ax.set_zlabel('Z') + ax.set_xlabel('X') + ax.set_ylabel('Y') + ax.set_zlabel('Z') ax.view_init(elev=12., azim=80) - plt.tick_params(left = False, right = False , labelleft = False , - labelbottom = False, bottom = False) + plt.tick_params(left = False, right = False , labelleft = True , + labelbottom = True, bottom = False) for i in range(len(joint_pairs)): limb = joint_pairs[i] xs, ys, zs = [np.array([j3d[limb[0], j], j3d[limb[1], j]]) for j in range(3)] diff --git a/train.py b/train.py index 5950ee6..b8ea49d 100644 --- a/train.py +++ b/train.py @@ -10,6 +10,7 @@ import copy import random import prettytable +import json import torch import torch.nn as nn @@ -26,6 +27,9 @@ from lib.data.datareader_h36m import DataReaderH36M from lib.model.loss import * +import os +os.environ['KMP_DUPLICATE_LIB_OK'] = 'TRUE' + def parse_args(): parser = argparse.ArgumentParser() parser.add_argument("--config", type=str, default="configs/pretrain.yaml", help="Path to the config file.") @@ -52,11 +56,13 @@ def save_checkpoint(chk_path, epoch, lr, optimizer, model_pos, min_loss): 'model_pos': model_pos.state_dict(), 'min_loss' : min_loss }, chk_path) - -def evaluate(args, model_pos, test_loader, datareader): + +def evaluate(args, model_pos, test_loader, datareader=None): + # i used Claude to generate a new evaluate function because the previous one was limited to just the original testing set print('INFO: Testing') results_all = [] - model_pos.eval() + gts_all = [] + model_pos.eval() with torch.no_grad(): for batch_input, batch_gt in tqdm(test_loader): N, T = batch_gt.shape[:2] @@ -64,91 +70,44 @@ def evaluate(args, model_pos, test_loader, datareader): batch_input = batch_input.cuda() if args.no_conf: batch_input = batch_input[:, :, :, :2] - if args.flip: + if args.flip: batch_input_flip = flip_data(batch_input) predicted_3d_pos_1 = model_pos(batch_input) predicted_3d_pos_flip = model_pos(batch_input_flip) - predicted_3d_pos_2 = flip_data(predicted_3d_pos_flip) # Flip back - predicted_3d_pos = (predicted_3d_pos_1+predicted_3d_pos_2) / 2 + predicted_3d_pos_2 = flip_data(predicted_3d_pos_flip) + predicted_3d_pos = (predicted_3d_pos_1 + predicted_3d_pos_2) / 2 else: predicted_3d_pos = model_pos(batch_input) + if args.rootrel: - predicted_3d_pos[:,:,0,:] = 0 # [N,T,17,3] + predicted_3d_pos[:, :, 0, :] = 0 else: - batch_gt[:,0,0,2] = 0 + batch_gt[:, 0, 0, 2] = 0 if args.gt_2d: - predicted_3d_pos[...,:2] = batch_input[...,:2] + predicted_3d_pos[..., :2] = batch_input[..., :2] + results_all.append(predicted_3d_pos.cpu().numpy()) - results_all = np.concatenate(results_all) - results_all = datareader.denormalize(results_all) - _, split_id_test = datareader.get_split_id() - actions = np.array(datareader.dt_dataset['test']['action']) - factors = np.array(datareader.dt_dataset['test']['2.5d_factor']) - gts = np.array(datareader.dt_dataset['test']['joints_2.5d_image']) - sources = np.array(datareader.dt_dataset['test']['source']) + gts_all.append(batch_gt.cpu().numpy()) - num_test_frames = len(actions) - frames = np.array(range(num_test_frames)) - action_clips = actions[split_id_test] - factor_clips = factors[split_id_test] - source_clips = sources[split_id_test] - frame_clips = frames[split_id_test] - gt_clips = gts[split_id_test] - assert len(results_all)==len(action_clips) - - e1_all = np.zeros(num_test_frames) - e2_all = np.zeros(num_test_frames) - oc = np.zeros(num_test_frames) - results = {} - results_procrustes = {} - action_names = sorted(set(datareader.dt_dataset['test']['action'])) - for action in action_names: - results[action] = [] - results_procrustes[action] = [] - block_list = ['s_09_act_05_subact_02', - 's_09_act_10_subact_02', - 's_09_act_13_subact_01'] - for idx in range(len(action_clips)): - source = source_clips[idx][0][:-6] - if source in block_list: - continue - frame_list = frame_clips[idx] - action = action_clips[idx][0] - factor = factor_clips[idx][:,None,None] - gt = gt_clips[idx] - pred = results_all[idx] - pred *= factor - - # Root-relative Errors - pred = pred - pred[:,0:1,:] - gt = gt - gt[:,0:1,:] - err1 = mpjpe(pred, gt) - err2 = p_mpjpe(pred, gt) - e1_all[frame_list] += err1 - e2_all[frame_list] += err2 - oc[frame_list] += 1 - for idx in range(num_test_frames): - if e1_all[idx] > 0: - err1 = e1_all[idx] / oc[idx] - err2 = e2_all[idx] / oc[idx] - action = actions[idx] - results[action].append(err1) - results_procrustes[action].append(err2) - final_result = [] - final_result_procrustes = [] - summary_table = prettytable.PrettyTable() - summary_table.field_names = ['test_name'] + action_names - for action in action_names: - final_result.append(np.mean(results[action])) - final_result_procrustes.append(np.mean(results_procrustes[action])) - summary_table.add_row(['P1'] + final_result) - summary_table.add_row(['P2'] + final_result_procrustes) - print(summary_table) - e1 = np.mean(np.array(final_result)) - e2 = np.mean(np.array(final_result_procrustes)) - print('Protocol #1 Error (MPJPE):', e1, 'mm') - print('Protocol #2 Error (P-MPJPE):', e2, 'mm') + results_all = np.concatenate(results_all) # (N_clips, T, 17, 3) + gts_all = np.concatenate(gts_all) # (N_clips, T, 17, 3) + + # Flatten clips into frames + N_clips, T, J, C = results_all.shape + pred_flat = results_all.reshape(-1, J, C) # (N_clips*T, 17, 3) + gt_flat = gts_all.reshape(-1, J, C) + + # Root-relative + pred_flat = pred_flat - pred_flat[:, 0:1, :] + gt_flat = gt_flat - gt_flat[:, 0:1, :] + + # Both functions expect numpy (N, J, C) and return (N,) + e1 = np.mean(mpjpe(pred_flat, gt_flat)) + e2 = np.mean(p_mpjpe(pred_flat, gt_flat)) + + print(f'Protocol #1 Error (MPJPE): {e1} mm') + print(f'Protocol #2 Error (P-MPJPE): {e2} mm') print('----------') return e1, e2, results_all @@ -219,19 +178,19 @@ def train_with_config(args, opts): trainloader_params = { 'batch_size': args.batch_size, 'shuffle': True, - 'num_workers': 12, + 'num_workers': 0, 'pin_memory': True, - 'prefetch_factor': 4, - 'persistent_workers': True + # 'prefetch_factor': 4, + # 'persistent_workers': True } testloader_params = { 'batch_size': args.batch_size, 'shuffle': False, - 'num_workers': 12, + 'num_workers': 0, 'pin_memory': True, - 'prefetch_factor': 4, - 'persistent_workers': True + # 'prefetch_factor': 4, + # 'persistent_workers': True } train_dataset = MotionDataset3D(args, args.subset_list, 'train') @@ -245,7 +204,7 @@ def train_with_config(args, opts): instav = InstaVDataset2D() instav_loader_2d = DataLoader(instav, **trainloader_params) - datareader = DataReaderH36M(n_frames=args.clip_len, sample_stride=args.sample_stride, data_stride_train=args.data_stride, data_stride_test=args.clip_len, dt_root = 'data/motion3d', dt_file=args.dt_file) + datareader = None min_loss = 100000 model_backbone = load_backbone(args) model_params = 0 @@ -284,6 +243,7 @@ def train_with_config(args, opts): if args.partial_train: model_pos = partial_train_layers(model_pos, args.partial_train) + full_stats = [] if not opts.evaluate: lr = args.learning_rate optimizer = optim.AdamW(filter(lambda p: p.requires_grad, model_pos.parameters()), lr=lr, weight_decay=args.weight_decay) @@ -330,12 +290,13 @@ def train_with_config(args, opts): train_epoch(args, model_pos, train_loader_3d, losses, optimizer, has_3d=True, has_gt=True) elapsed = (time() - start_time) / 60 - if args.no_eval: + if args.no_eval or (args.eval_last and epoch != args.epochs - 1): print('[%d] time %.2f lr %f 3d_train %f' % ( epoch + 1, elapsed, lr, losses['3d_pos'].avg)) + e1, e2 = min_loss, min_loss else: e1, e2, results_all = evaluate(args, model_pos, test_loader, datareader) print('[%d] time %.2f lr %f 3d_train %f e1 %f e2 %f' % ( @@ -355,6 +316,18 @@ def train_with_config(args, opts): train_writer.add_scalar('loss_a', losses['angle'].avg, epoch + 1) train_writer.add_scalar('loss_av', losses['angle_velocity'].avg, epoch + 1) train_writer.add_scalar('loss_total', losses['total'].avg, epoch + 1) + + full_stats.append({ + "loss_3d_pos": losses['3d_pos'].avg, + "loss_2d_proj": losses['2d_proj'].avg, + "loss_3d_scale": losses['3d_scale'].avg, + "loss_3d_velocity": losses['3d_velocity'].avg, + "loss_lv": losses['lv'].avg, + "loss_lg": losses['lg'].avg, + "loss_a": losses['angle'].avg, + "loss_av": losses['angle_velocity'].avg, + "loss_total": losses['total'].avg + }) # Decay learning rate exponentially lr *= lr_decay @@ -375,6 +348,10 @@ def train_with_config(args, opts): if opts.evaluate: e1, e2, results_all = evaluate(args, model_pos, test_loader, datareader) + + res_path = os.path.join(opts.checkpoint, 'results.json') + with open(res_path, 'w') as f: + json.dump(full_stats, f) if __name__ == "__main__": opts = parse_args()