diff --git a/configs/mvhn_100840/avatar.yaml b/configs/mvhn_100840/avatar.yaml new file mode 100644 index 0000000..70c3c9d --- /dev/null +++ b/configs/mvhn_100840/avatar.yaml @@ -0,0 +1,75 @@ +mode: train +train: + dataset: MvRgbDatasetMVHumanNet + data: + subject_name: 100840 + data_dir: ../MVHumanNet_plusplus/mvhumannet++_10demo/main/100840 + frame_range: &id001 + - 0 + - 60 + - 1 + used_cam_ids: + - 0 + - 1 + - 2 + - 3 + - 4 + - 5 + - 6 + - 8 + - 9 + - 10 + - 11 + - 12 + - 14 + - 15 + load_smpl_pos_map: true + pretrained_dir: null + net_ckpt_dir: ./results/100840/avatar + prev_ckpt: null + ckpt_interval: + epoch: 10 + batch: 50000 + eval_interval: 1000 + eval_training_ids: + - 42 + - 2 + eval_testing_ids: + - 50 + - 0 + eval_img_factor: 1.0 + lr_init: 0.0005 + loss_weight: + l1: 1.0 + lpips: 0.1 + offset: 0.005 + finetune_color: false + batch_size: 1 + num_workers: 8 + random_bg_color: true +test: + dataset: MvRgbDatasetMVHumanNet + data: + data_dir: ../MVHumanNet_plusplus/mvhumannet++_10demo/main/100840 + frame_range: [0, 60] + subject_name: avatarrex_zzr + pose_data: + data_path: ../thuman4/pose_00.npz + frame_range: [2000, 2500] +# data_path: Z:/Data/Pose/AMASS/CMU/06/06_13_poses.npz +# data_path: Z:/Data/Pose/AMASS/CMU/10/10_05_poses.npz +# frame_interval: 4 + view_setting: free + render_view_idx: 13 + global_orient: true + img_scale: 1.0 + save_mesh: false + render_skeleton: false + save_tex_map: false + save_ply: false + n_pca: 20 + sigma_pca: 2.0 + prev_ckpt: None +model: + with_viewdirs: true + random_style: false diff --git a/configs/mvhn_100840/template.yaml b/configs/mvhn_100840/template.yaml new file mode 100644 index 0000000..1d94ec3 --- /dev/null +++ b/configs/mvhn_100840/template.yaml @@ -0,0 +1,46 @@ +mode: train +train: + dataset: MvRgbDatasetMVHumanNet + data: + subject_name: 100840 + data_dir: ../MVHumanNet_plusplus/mvhumannet++_10demo/main/100840 + frame_range: &id001 + - 60 + - 61 + - 1 + used_cam_ids: null + mode: nerf + net_ckpt_dir: ./results/100840/template + prev_ckpt: null + save_init_ckpt: false + start_epoch: 0 + end_epoch: 1001 + ckpt_interval: + epoch: 500 + batch: 1000 + eval_interval: 10000 + depth_guided_sampling: + flag: false + near_sur_dist: 0.05 + N_ray_samples: 32 + compute_grad: true + lr: + network: + type: Step + initial: 0.0005 + interval: 40000 + factor: 0.9 + min: 8.0e-05 + loss_weight: + color: 1.0 + mask: 1.0 + eikonal: 0.1 + batch_size: 1 + num_workers: 0 +model: + with_hand: true + volume_type: diff + use_root_finding: true + multires: 6 + use_viewdir: false + multires_viewdir: 3 diff --git a/configs/mvhn_100990/avatar.yaml b/configs/mvhn_100990/avatar.yaml new file mode 100644 index 0000000..803bf80 --- /dev/null +++ b/configs/mvhn_100990/avatar.yaml @@ -0,0 +1,75 @@ +mode: train +train: + dataset: MvRgbDatasetMVHumanNet + data: + subject_name: 100990 + data_dir: ../MVHumanNet_plusplus/test/100990 + frame_range: &id001 + - 0 + - 60 + - 1 + used_cam_ids: + - 0 + - 1 + - 2 + - 3 + - 4 + - 5 + - 6 + - 8 + - 9 + - 10 + - 11 + - 12 + - 14 + - 15 + load_smpl_pos_map: true + pretrained_dir: ./results/100990/avatar/pretrained + net_ckpt_dir: ./results/100990/avatar + prev_ckpt: null + ckpt_interval: + epoch: 30 + batch: 50000 + eval_interval: 1000 + eval_training_ids: + - 1 + - 1 + eval_testing_ids: + - 9 + - 14 + eval_img_factor: 1.0 + lr_init: 0.0005 + loss_weight: + l1: 1.0 + lpips: 0.1 + offset: 0.005 + finetune_color: false + batch_size: 1 + num_workers: 8 + random_bg_color: true +test: + dataset: MvRgbDatasetMVHumanNet + data: + data_dir: ../MVHumanNet_plusplus/test/100990 + frame_range: [0, 60] + subject_name: avatarrex_zzr + pose_data: + data_path: ../thuman4/pose_00.npz + frame_range: [2000, 2500] +# data_path: Z:/Data/Pose/AMASS/CMU/06/06_13_poses.npz +# data_path: Z:/Data/Pose/AMASS/CMU/10/10_05_poses.npz +# frame_interval: 4 + view_setting: free + render_view_idx: 13 + global_orient: true + img_scale: 1.0 + save_mesh: false + render_skeleton: false + save_tex_map: false + save_ply: false + n_pca: 20 + sigma_pca: 2.0 + prev_ckpt: None +model: + with_viewdirs: true + random_style: false diff --git a/configs/mvhn_100990/template.yaml b/configs/mvhn_100990/template.yaml new file mode 100644 index 0000000..e9f03dd --- /dev/null +++ b/configs/mvhn_100990/template.yaml @@ -0,0 +1,46 @@ +mode: train +train: + dataset: MvRgbDatasetMVHumanNet + data: + subject_name: 100990 + data_dir: ../MVHumanNet_plusplus/test/100990 + frame_range: &id001 + - 60 + - 61 + - 1 + used_cam_ids: null + mode: nerf + net_ckpt_dir: ./results/100990/template + prev_ckpt: null + save_init_ckpt: false + start_epoch: 0 + end_epoch: 1001 + ckpt_interval: + epoch: 500 + batch: 1000 + eval_interval: 10000 + depth_guided_sampling: + flag: false + near_sur_dist: 0.05 + N_ray_samples: 32 + compute_grad: true + lr: + network: + type: Step + initial: 0.0005 + interval: 40000 + factor: 0.9 + min: 8.0e-05 + loss_weight: + color: 1.0 + mask: 1.0 + eikonal: 0.1 + batch_size: 1 + num_workers: 0 +model: + with_hand: true + volume_type: diff + use_root_finding: true + multires: 6 + use_viewdir: false + multires_viewdir: 3 diff --git a/dataset/dataset_mv_rgb.py b/dataset/dataset_mv_rgb.py index 439dca9..e78c096 100644 --- a/dataset/dataset_mv_rgb.py +++ b/dataset/dataset_mv_rgb.py @@ -498,3 +498,50 @@ def load_color_mask_images(self, pose_idx, view_idx): color_img = cv.imread(self.data_dir + '/4x/rgbs/%s/%s_rgb%06d.jpg' % (cam_name, cam_name, pose_idx), cv.IMREAD_UNCHANGED) mask_img = cv.imread(self.data_dir + '/4x/masks/%s/%s_mask%06d.png' % (cam_name, cam_name, pose_idx), cv.IMREAD_UNCHANGED) return color_img, mask_img + +## -------------------------------------------------- ## +## ------- NEW CLASS FOR MVHumanNet DATASET --------- ## +class MvRgbDatasetMVHumanNet(MvRgbDatasetBase): + def __init__( + self, + data_dir, + frame_range = None, + used_cam_ids = None, + training = True, + subject_name = None, + load_smpl_pos_map = False, + load_smpl_nml_map = False, + mode = '3dgs' + ): + super(MvRgbDatasetMVHumanNet, self).__init__( + data_dir, + frame_range, + used_cam_ids, + training, + subject_name, + load_smpl_pos_map, + load_smpl_nml_map, + mode + ) + self.img_files = sorted(os.listdir(self.data_dir + '/images/cam_00')) + self.img_file_names = [f.rstrip('.jpg') for f in self.img_files] + + def load_cam_data(self): + import json + cam_data = json.load(open(self.data_dir + '/calibrations.json', 'r')) + self.view_num = len(cam_data) + self.extr_mats = [] + self.cam_names = ['cam%02d' % view_idx for view_idx in range(self.view_num)] + for view_idx in range(self.view_num): + extr_mat = np.identity(4, np.float32) + extr_mat[:3, :3] = np.array(cam_data['cam_%02d' % view_idx]['R'], np.float32).reshape(3, 3) + extr_mat[:3, 3] = np.array(cam_data['cam_%02d' % view_idx]['T'], np.float32) + self.extr_mats.append(extr_mat) + self.intr_mats = [np.array(cam_data['cam_%02d' % view_idx]['K'], np.float32).reshape(3, 3) for view_idx in range(self.view_num)] + self.img_heights = [cam_data['cam_%02d' % view_idx]['imgSize'][1] for view_idx in range(self.view_num)] + self.img_widths = [cam_data['cam_%02d' % view_idx]['imgSize'][0] for view_idx in range(self.view_num)] + + def load_color_mask_images(self, pose_idx, view_idx): + color_img = cv.imread(f'{self.data_dir}/images/cam_{view_idx:02d}/{self.img_file_names[pose_idx]}.jpg', cv.IMREAD_UNCHANGED) + mask_img = cv.imread(f'{self.data_dir}/masks/cam_{view_idx:02d}/{self.img_file_names[pose_idx]}.png', cv.IMREAD_UNCHANGED) + return color_img, mask_img diff --git a/gen_data/TRAIN_MVHUMANNET++_DATASET.md b/gen_data/TRAIN_MVHUMANNET++_DATASET.md new file mode 100644 index 0000000..354ccae --- /dev/null +++ b/gen_data/TRAIN_MVHUMANNET++_DATASET.md @@ -0,0 +1,45 @@ +# Train Models Using the MVHumanNet++ Dataset + +## MVHumanNet++ Dataset +Download MVHumanNet++ Dataset from the [official repo](https://github.com/GAP-LAB-CUHK-SZ/MVHumanNet_plusplus). + +MVHumanNet++ is very nice for training AnimatableGaussians models. It features: +- a large number of subjects and outfits +- 16 cameras surrounding the subject +- 60 frames for each subject-outfit-camera and **an extra A-pose frame** +- camera calibrations, smplx-parameters are accurate and complete + +## Preporcessing +If you want to preprocess one single subject, run +```bash +python preprocess_mvhnpp.py --mvhn_subj_path ../../mvhumannet_plusplus/test/100990 +``` + +You can put multiple subjects in a folder (e.g. `../../mvhumannet_plusplus/test/`) and preprocess them together by running: +```bash +python preprocess_mvhnpp.py --mvhn_root ../../mvhumannet_plusplus/test +``` + +The script will create `calibrations.json` and `smpl_params.npz` in the subject folder. These files conform to the THuman4 format, as the authors of AnimatableGaussians recommended. + +## Training + +You can follow [GEN_DATA.md](./GEN_DATA.md) and [README.md](../README.md) to train or directly use the script. From the root of AnimatableGaussians +```bash +chmod +x train_mvhnpp.sh +./train_mvhnpp.sh mvhn_100990 +``` +You must have `template.yaml` and `avatar.yaml` in `./configs/mvhn_100990` for this to work, which we have provided as an example. + +### P.S. +I could not configure the env using the given `requirements.txt`. I am using +- `python=3.10` +- `cuda=11.7.1` +- `pytorch-cuda=2.0.1` +- `numpy=1.26.4` +- `pytorch3d=0.7.4` (installed using downloaded pre-compiled file) +- **`scikit-learn=1.1.3`** (different from `requirements.txt`) +- **`scikit-image=0.25.2`** (different from `requirements.txt`) +- **`opencv=4.7.0`** (different from `requirements.txt`) +- **`setuptools=78.1.0`** (different from `requirements.txt`) +- `smplx=0.1.28` \ No newline at end of file diff --git a/gen_data/preprocess_mvhnpp.py b/gen_data/preprocess_mvhnpp.py new file mode 100644 index 0000000..15ea6de --- /dev/null +++ b/gen_data/preprocess_mvhnpp.py @@ -0,0 +1,109 @@ +import os +from pathlib import Path +import json +import numpy as np +import smplx +import torch + +def conv_mvhumannetpp_cameras_to_thuman(mvhn_subj_path): + cameras = sorted(os.listdir(mvhn_subj_path / 'cameras')) + calibrations = {} + for cam in cameras: + calibration = np.load(mvhumannet_root / subj / 'cameras' / cam / 'camera.npz') + K = np.array(calibration['intrinsic']).flatten().tolist() + ext = calibration['extrinsic'] + R = np.array(ext[:3, :3]).flatten().tolist() + T = np.array(ext[:3, 3]).tolist() + cam_cal = {'K': K, 'R': R, 'T': T, + 'distCoeff': [0.0, 0.0, 0.0, 0.0], + 'imgSize': [1024, 1024], + 'rectifyAlpha': 0.0} + calibrations[cam] = cam_cal + with open(mvhn_subj_path / 'calibrations.json', 'w') as f: + json.dump(calibrations, f) + + +def convert_hand_pca_to_rotation(pca_params, side='right'): + """ + Converts 6-dim PCA parameters to 45-dim Axis-Angle parameters. + + Args: + pca_params (numpy array): Shape (1, 6) or (N, 6) + side (str): 'left' or 'right' + + Returns: + numpy array: Shape (1, 45) + """ + # --- CONFIGURATION --- + MODEL_FOLDER = '../smpl_files' # Path to folder containing SMPLX_NEUTRAL.npz + GENDER = 'neutral' + + model = smplx.create( + MODEL_FOLDER, + model_type='smplx', + gender=GENDER, + use_pca=True, + num_pca_comps=6 + ) + + # Convert input to tensor + pca_tensor = torch.tensor(pca_params, dtype=torch.float32) + + # 2. Extract the specific data for the requested hand + if side == 'right': + mean_pose = model.right_hand_mean + components = model.right_hand_components + else: + mean_pose = model.left_hand_mean + components = model.left_hand_components + + active_components = components[:6, :] + delta = torch.matmul(pca_tensor, active_components) + full_pose = delta + mean_pose + + return full_pose.detach().numpy() + + +def conv_mvhumannetpp_smplx_to_thuman(mvhn_subj_path): + frames = sorted(os.listdir(mvhn_subj_path / 'smplx_params_new')) + smplx_all_frames = [np.load(mvhn_subj_path / 'smplx_params_new' / fr) for fr in frames] + betas = smplx_all_frames[0]['betas'] + transl = np.concatenate([smplx_fr['transl'] for smplx_fr in smplx_all_frames]) + global_orient = np.concatenate([smplx_fr['global_orient'] for smplx_fr in smplx_all_frames]) + body_pose = np.concatenate([smplx_fr['body_pose'] for smplx_fr in smplx_all_frames]) + jaw_pose = np.concatenate([smplx_fr['jaw_pose'] for smplx_fr in smplx_all_frames]) + expression = np.concatenate([smplx_fr['expression'] for smplx_fr in smplx_all_frames]) + left_hand_pose = np.concatenate([convert_hand_pca_to_rotation(smplx_fr['left_hand_pose'], 'left') + for smplx_fr in smplx_all_frames]) + right_hand_pose = np.concatenate([convert_hand_pca_to_rotation(smplx_fr['right_hand_pose'], 'right') + for smplx_fr in smplx_all_frames]) + + np.savez(mvhn_subj_path / 'smpl_params.npz', + betas=betas, + global_orient=global_orient, + transl=transl, + body_pose=body_pose, + jaw_pose=jaw_pose, + expression=expression, + left_hand_pose=left_hand_pose, + right_hand_pose=right_hand_pose) + + +if __name__ == '__main__': + from argparse import ArgumentParser + + arg_parser = ArgumentParser() + arg_parser.add_argument('--mvhn_subj_path', type = str, help = 'Path to a single MVHumanNet++ subject.') + arg_parser.add_argument('--mvhn_root', type = str, help = 'Path to folder containing multiple subjects.') + args = arg_parser.parse_args() + + if args.mvhn_subj_path is not None: + mvhn_subj_paths = [Path(args.mvhn_subj_path)] + else: + mvhn_root = Path(args.mvhn_root) + subjects = sorted(os.listdir(mvhn_root)) + mvhn_subj_paths = [mvhn_root / subj for subj in subjects] + + for mvhn_subj_path in mvhn_subj_paths: + conv_mvhumannetpp_cameras_to_thuman(mvhn_subj_path) + conv_mvhumannetpp_smplx_to_thuman(mvhn_subj_path) \ No newline at end of file diff --git a/train_mvhnpp.sh b/train_mvhnpp.sh new file mode 100644 index 0000000..c9a474f --- /dev/null +++ b/train_mvhnpp.sh @@ -0,0 +1,22 @@ +#!/bin/bash + +# Check if an argument was provided +if [ -z "$1" ]; then + echo "Error: No config directory provided." + echo "Usage: ./run_gen.sh " + exit 1 +fi + +SUBJECT_NAME="$1" + +echo "Generating canonical LBS weight volume ..." +python -m gen_data.gen_weight_volume -c "configs/${SUBJECT_NAME}/template.yaml" + +echo "Reconstructing a template ..." +python main_template.py -c "configs/${SUBJECT_NAME}/template.yaml" + +echo "Generating position maps ..." +OPENCV_IO_ENABLE_OPENEXR=1 python -m gen_data.gen_pos_maps -c "configs/${SUBJECT_NAME}/avatar.yaml" + +echo "Traning avatar ..." +python main_avatar.py -c "configs/${SUBJECT_NAME}/avatar.yaml" --mode=train