diff --git a/.gitignore b/.gitignore index b6e4761..3838b9d 100644 --- a/.gitignore +++ b/.gitignore @@ -127,3 +127,14 @@ dmypy.json # Pyre type checker .pyre/ + +exps_first_stage/ +exps_second_stage/ +implicit-hair-data/ +pretrained_models/ + +npbgpp/build/ +*/__pycache__/ + +src/multiview_optimization/experiments/ +src/multiview_optimization/PIXIE/ diff --git a/CDGNet b/CDGNet index 9daf7dd..9d7168e 160000 --- a/CDGNet +++ b/CDGNet @@ -1 +1 @@ -Subproject commit 9daf7ddee6045c151c90a2e300946ea5f5717591 +Subproject commit 9d7168e457dc55a12edcaf4eafbeaf2d3bd8f25d diff --git a/neural_haircut.yaml b/neural_haircut.yaml index a10496e..bb419a0 100644 --- a/neural_haircut.yaml +++ b/neural_haircut.yaml @@ -162,7 +162,7 @@ dependencies: - ncurses=6.4=h6a678d5_0 - nest-asyncio=1.5.6=py39h06a4308_0 - nettle=3.7.3=hbbd107a_1 - - notebook=6.5.3=py39h06a4308_0 + - notebook=6.5.* - notebook-shim=0.2.2=py39h06a4308_0 - nsight-compute=2023.1.0.15=0 - nspr=4.33=h295c915_0 @@ -285,6 +285,7 @@ dependencies: - icecream==2.1.3 - imageio==2.27.0 - importlib-resources==5.12.0 + - joblib==1.3.2 - jsonmerge==1.9.0 - k-diffusion==0.0.12 - kiwisolver==1.4.4 @@ -306,7 +307,7 @@ dependencies: - pathspec==0.11.1 - pathtools==0.1.2 - plotly==5.14.0 - - protobuf==4.22.1 + - protobuf==3.19.6 - pyasn1==0.4.8 - pyasn1-modules==0.2.8 - pycodestyle==2.10.0 diff --git a/preprocess_custom_data/calc_masks.py b/preprocess_custom_data/calc_masks.py index 4cbbbd6..01986d3 100644 --- a/preprocess_custom_data/calc_masks.py +++ b/preprocess_custom_data/calc_masks.py @@ -124,6 +124,9 @@ def main(args): images = sorted(os.listdir(os.path.join(args.scene_path, 'image'))) n_images = len(sorted(os.listdir(os.path.join(args.scene_path, 'image')))) + + # Get image file extension + image_ext = os.path.splitext(images[0])[1] tens_list = [] for i in range(n_images): @@ -156,7 +159,8 @@ def main(args): state_dict = model.state_dict().copy() state_dict_old = torch.load(args.CDGNET_ckpt, map_location='cpu') - for key, nkey in zip(state_dict_old.keys(), state_dict.keys()): + state_dict_keys = list(state_dict.keys()) + for key, nkey in zip(state_dict_old.keys(), state_dict_keys): if key != nkey: # remove the 'module.' in the 'key' state_dict[key[7:]] = deepcopy(state_dict_old[key]) @@ -174,11 +178,11 @@ def main(args): images = [] masks = [] for basename in basenames: - img = Image.open(os.path.join(args.scene_path, 'image', basename + '.jpg')) + img = Image.open(os.path.join(args.scene_path, 'image', basename + image_ext)) raw_images.append(np.asarray(img)) img = transform(img.resize(input_size))[None] img = torch.cat([img, torch.flip(img, dims=[-1])], dim=0) - mask = np.asarray(Image.open(os.path.join(args.scene_path, 'mask', basename + '.jpg'))) + mask = np.asarray(Image.open(os.path.join(args.scene_path, 'mask', basename + '.png'))) images.append(img) masks.append(mask) @@ -188,7 +192,7 @@ def main(args): for i in range(len(images)): hair_mask = np.asarray(Image.fromarray((parsing_preds[i] == 2)).resize(image_size, Image.BICUBIC)) hair_mask = hair_mask * masks[i] - Image.fromarray(hair_mask).save(os.path.join(args.scene_path, 'hair_mask', basenames[i] + '.jpg')) + Image.fromarray(hair_mask).save(os.path.join(args.scene_path, 'hair_mask', basenames[i] + '.png')) print('Results saved in folder: ', os.path.join(args.scene_path, 'hair_mask')) @@ -197,7 +201,7 @@ def main(args): parser.add_argument('--scene_path', default='./implicit-hair-data/data/h3ds/168f8ca5c2dce5bc/', type=str) parser.add_argument('--MODNET_ckpt', default='./MODNet/pretrained/modnet_photographic_portrait_matting.ckpt', type=str) - parser.add_argument('--CDGNET_ckpt', default='./cdgnet/snapshots/LIP_epoch_149.pth', type=str) + parser.add_argument('--CDGNET_ckpt', default='./CDGNet/snapshots/LIP_epoch_149.pth', type=str) args, _ = parser.parse_known_args() args = parser.parse_args() diff --git a/preprocess_custom_data/calc_orientation_maps.py b/preprocess_custom_data/calc_orientation_maps.py index 3e1977b..55ba154 100644 --- a/preprocess_custom_data/calc_orientation_maps.py +++ b/preprocess_custom_data/calc_orientation_maps.py @@ -8,6 +8,7 @@ import tqdm import cv2 import argparse +from joblib import Parallel, delayed def rgb2gray(rgb): r, g, b = rgb[:,:,0], rgb[:,:,1], rgb[:,:,2] @@ -25,23 +26,23 @@ def generate_gabor_filters(sigma_x, sigma_y, freq, num_filters): return kernels -def calc_orients(img, kernels): +def calc_orients(img, kernels, n_jobs): gray_img = rgb2gray(img) filtered_image = difference_of_gaussians(gray_img, 0.4, 10) - gabor_filtered_images = [ndi.convolve(filtered_image, kernels[i], mode='wrap') for i in range(len(kernels))] + gabor_filtered_images = list(Parallel(n_jobs=n_jobs, prefer="threads")(delayed(ndi.convolve)(filtered_image, kernels[i], mode='wrap') for i in range(len(kernels)))) F_orients = np.abs(np.stack(gabor_filtered_images)) # abs because we only measure angle in [0, pi] return F_orients -def calc_confidences(F_orients, orientation_map): +def calc_confidences(F_orients, orientation_map, num_filters, n_jobs): orients_bins = np.linspace(0, math.pi * (num_filters - 1) / num_filters, num_filters) - orients_bins = orients_bins[:, None, None] - orientation_map = orientation_map[None] - - dists = np.minimum(np.abs(orientation_map - orients_bins), - np.minimum(np.abs(orientation_map - orients_bins - math.pi), - np.abs(orientation_map - orients_bins + math.pi))) + def calc_dists(i): + dists = np.abs(orientation_map - orients_bins[i]) + dists = np.minimum(np.abs(orientation_map - orients_bins[i] - math.pi), dists) + dists = np.minimum(np.abs(orientation_map - orients_bins[i] + math.pi), dists) + return dists + dists = np.stack(Parallel(n_jobs=n_jobs, prefer="threads")(delayed(calc_dists)(i) for i in range(num_filters))) F_orients_norm = F_orients / F_orients.sum(axis=0, keepdims=True) @@ -56,17 +57,17 @@ def main(args): kernels = generate_gabor_filters(args.sigma_x, args.sigma_y, args.freq, args.num_filters) - img_list = sorted(os.listdir(img_path)) + img_list = sorted(os.listdir(args.img_path)) for img_name in tqdm.tqdm(img_list): basename = img_name.split('.')[0] - img = np.array(Image.open(os.path.join(img_path, img_name))) - F_orients = calc_orients(img, kernels) + img = np.array(Image.open(os.path.join(args.img_path, img_name))) + F_orients = calc_orients(img, kernels, args.n_jobs) orientation_map = F_orients.argmax(0) - orientation_map_rad = orientation_map.astype('float16') / num_filters * math.pi - confidence_map = calc_confidences(F_orients, orientation_map_rad) + orientation_map_rad = orientation_map.astype('float16') / args.num_filters * math.pi + confidence_map = calc_confidences(F_orients, orientation_map_rad, args.num_filters, args.n_jobs) - cv2.imwrite(f'{orient_dir}/{basename}.png', orientation_map.astype('uint8')) - np.save(f'{conf_dir}/{basename}.npy', confidence_map.astype('float16')) + cv2.imwrite(f'{args.orient_dir}/{basename}.png', orientation_map.astype('uint8')) + np.save(f'{args.conf_dir}/{basename}.npy', confidence_map.astype('float16')) @@ -80,6 +81,7 @@ def main(args): parser.add_argument('--sigma_y', default=2.4, type=float) parser.add_argument('--freq', default=0.23, type=float) parser.add_argument('--num_filters', default=180, type=int) + parser.add_argument('--n_jobs', default=8, type=int) args, _ = parser.parse_known_args() args = parser.parse_args() diff --git a/preprocess_custom_data/colmap_parsing.py b/preprocess_custom_data/colmap_parsing.py index dedb50e..b92ee48 100644 --- a/preprocess_custom_data/colmap_parsing.py +++ b/preprocess_custom_data/colmap_parsing.py @@ -11,9 +11,9 @@ def main(args): - images_file = f'{args.path_to_scene}/sparse_txt/images.txt' - points_file = f'{args.path_to_scene}/sparse_txt/points3D.txt' - camera_file = f'{args.path_to_scene}/sparse_txt/cameras.txt' + images_file = f'{args.path_to_scene}/colmap/sparse_txt/images.txt' + points_file = f'{args.path_to_scene}/colmap/sparse_txt/points3D.txt' + camera_file = f'{args.path_to_scene}/colmap/sparse_txt/cameras.txt' # Parse colmap cameras and used images with open(camera_file) as f: @@ -37,18 +37,15 @@ def main(args): data = {} - image_names = [] for i in range(n_images): line_split = images_file_lines[4 + i * 2].split() - image_id = int(line_split[0]) q = np.array([float(x) for x in line_split[1: 5]]) # w, x, y, z t = np.array([float(x) for x in line_split[5: 8]]) image_name = line_split[-1] - image_names.append(image_name) extrinsic_matrix = np.eye(4) extrinsic_matrix[:3, :3] = R.from_quat(np.roll(q, -1)).as_matrix() @@ -73,31 +70,20 @@ def main(args): points = np.stack(points) colors = np.stack(colors) - output_folder = args.save_path - images_folder = os.path.join(args.path_to_scene, 'video_frames') - - os.makedirs(output_folder, exist_ok=True) - os.makedirs(os.path.join(output_folder, 'full_res_image'), exist_ok=True) - cameras = [] - debug = False - for i, k in enumerate(data.keys()): - filename = f'img_{i:04}.png' + for k in sorted(data.keys()): T = data[k] cameras.append(T) - shutil.copyfile(os.path.join(images_folder, k), os.path.join(output_folder, 'full_res_image', filename)) - np.savez(os.path.join(output_folder, 'cameras.npz'), np.stack(cameras)) - trimesh.points.PointCloud(points).export(os.path.join(output_folder, 'point_cloud.ply')); + np.savez(os.path.join(args.path_to_scene, 'cameras.npz'), np.stack(cameras)) + trimesh.points.PointCloud(points).export(os.path.join(args.path_to_scene, 'point_cloud.ply')); if __name__ == "__main__": parser = argparse.ArgumentParser(conflict_handler='resolve') parser.add_argument('--path_to_scene', default='./implicit-hair-data/data/monocular/person_1', type=str) - parser.add_argument('--save_path', default='./implicit-hair-data/data/monocular/person_1/colmap', type=str) - args, _ = parser.parse_known_args() args = parser.parse_args() diff --git a/preprocess_custom_data/copy_checkpoints.py b/preprocess_custom_data/copy_checkpoints.py index b37eb9b..935e00c 100644 --- a/preprocess_custom_data/copy_checkpoints.py +++ b/preprocess_custom_data/copy_checkpoints.py @@ -1,6 +1,7 @@ from shutil import copyfile import os from pathlib import Path +import argparse def main(args): @@ -11,6 +12,7 @@ def main(args): exps_dir = Path('./exps_first_stage') / exp_name / case / Path(conf_path).stem prev_exps = sorted(exps_dir.iterdir()) cur_dir = prev_exps[-1].name + scene_type = Path(conf_path).parent.stem path_to_mesh = os.path.join(exps_dir, cur_dir, 'meshes') path_to_ckpt = os.path.join(exps_dir, cur_dir, 'checkpoints') @@ -19,19 +21,19 @@ def main(args): meshes = sorted(os.listdir(path_to_mesh)) last_ckpt = sorted(os.listdir(path_to_ckpt))[-1] - last_hair = [i for i in head_string if i.split('_')[-1].split('.')[0]=='hair'][-1] - last_head = [i for i in head_string if i.split('_')[-1].split('.')[0]=='head'][-1] + last_hair = [i for i in meshes if i.split('_')[-1].split('.')[0]=='hair'][-1] + last_head = [i for i in meshes if i.split('_')[-1].split('.')[0]=='head'][-1] - print(f'Copy obtained from the first stage checkpoint, hair and head geometry to folder ./implicit-hair-data/data/{case}') + print(f'Copy obtained from the first stage checkpoint, hair and head geometry to folder ./implicit-hair-data/data/{scene_type}/{case}') - copyfile(os.path.join(path_to_mesh, last_hair), f'./implicit-hair-data/data/{case}/final_hair.ply') - copyfile(os.path.join(path_to_mesh, last_head), f'./implicit-hair-data/data/{case}/final_head.ply') - copyfile(os.path.join(path_to_ckpt, last_ckpt), f'./implicit-hair-data/data/{case}/ckpt_final.ply') + copyfile(os.path.join(path_to_mesh, last_hair), f'./implicit-hair-data/data/{scene_type}/{case}/final_hair.ply') + copyfile(os.path.join(path_to_mesh, last_head), f'./implicit-hair-data/data/{scene_type}/{case}/final_head.ply') + copyfile(os.path.join(path_to_ckpt, last_ckpt), f'./implicit-hair-data/data/{scene_type}/{case}/ckpt_final.pth') if os.path.exists(path_to_fitted_camera): - print(f'Copy obtained from the first stage camera fitting checkpoint to folder ./implicit-hair-data/data/{case}') + print(f'Copy obtained from the first stage camera fitting checkpoint to folder ./implicit-hair-data/data/{scene_type}/{case}') last_camera = sorted(os.listdir(path_to_fitted_camera))[-1] - copyfile(os.path.join(path_to_fitted_camera, last_camera), f'./implicit-hair-data/data/{case}/fitted_cameras.pth') + copyfile(os.path.join(path_to_fitted_camera, last_camera), f'./implicit-hair-data/data/{scene_type}/{case}/fitted_cameras.pth') diff --git a/preprocess_custom_data/cut_eyes.py b/preprocess_custom_data/cut_eyes.py index f166b3b..1480ac5 100644 --- a/preprocess_custom_data/cut_eyes.py +++ b/preprocess_custom_data/cut_eyes.py @@ -2,6 +2,7 @@ import os from pytorch3d.io import load_obj, save_obj import torch +import argparse def main(args): diff --git a/preprocess_custom_data/cut_scalp.py b/preprocess_custom_data/cut_scalp.py index 85f1e9f..54a06d2 100644 --- a/preprocess_custom_data/cut_scalp.py +++ b/preprocess_custom_data/cut_scalp.py @@ -90,7 +90,7 @@ def main(args): if face[0] in full_scalp_list and face[1] in full_scalp_list and face[2] in full_scalp_list: faces_masked.append(torch.tensor([d[int(face[0])], d[int(face[1])], d[int(face[2])]])) # print(faces_masked, full_scalp_list) - save_obj(os.path.join(save_path, 'scalp.obj'), scalp_mesh.verts_packed()[full_scalp_list], torch.stack(faces_masked)) + save_obj(os.path.join(save_path, 'scalp.obj'), scalp_mesh.verts_packed()[full_scalp_list], torch.stack(faces_masked)) with open(os.path.join(save_path, 'cut_scalp_verts.pickle'), 'wb') as f: pickle.dump(list(torch.tensor(sorted_idx).detach().cpu().numpy()), f) diff --git a/preprocess_custom_data/readme.md b/preprocess_custom_data/readme.md index b932ef4..c56a326 100644 --- a/preprocess_custom_data/readme.md +++ b/preprocess_custom_data/readme.md @@ -11,15 +11,11 @@ The full data folder is organized as follows: |-- case_name |-- video_frames # after parsing .mp4 (optional) |-- colmap # (optional) - |-- full_res_image - |-- cameras.npz - |-- point_cloud.ply |-- database.db |-- sparse - |-- dense |-- sparse_txt - |-- cameras.npz # camera parameters + |-- point_cloud.ply |-- image |-- mask |-- hair_mask @@ -47,7 +43,7 @@ The full data folder is organized as follows: ##### Run commands ```bash -colmap automatic_reconstructor --workspace_path CASE_NAME/colmap --image_path CASE_NAME/video_frames +colmap automatic_reconstructor --workspace_path CASE_NAME/colmap --image_path CASE_NAME/image --single_camera 1 --dense 0 ``` ```bash @@ -59,16 +55,17 @@ mkdir CASE_NAME/colmap/sparse_txt && colmap model_converter --input_path CASE_NA ##### To postprocess COLMAP's output run: ```bash -python colmap_parsing.py --path_to_scene ./implicit-hair-data/data/SCENE_TYPE/CASE --save_path ./implicit-hair-data/data/SCENE_TYPE/CASE/colmap +python preprocess_custom_data/colmap_parsing.py --path_to_scene ./implicit-hair-data/data/SCENE_TYPE/CASE ``` ##### Obtain: -After this step you would obtain ```colmap/full_res_image, colmap/cameras.npz, colmap/point_cloud.ply``` +After this step you would obtain ```cameras.npz``` and ```point_cloud.ply``` in ```./implicit-hair-data/data/SCENE_TYPE/CASE```. +Optionally you can run `verify_camera.py` to confirm that the camera parameters are correctly set. #### Step 2. (Optional) Define the region of interests in obtained point cloud. -Obtained ```colmap/point_cloud.ply``` is very noisy, so we are additionally define the region of interest using MeshLab and upload it to the current folder as ```point_cloud_cropped.ply```. +Obtained ```point_cloud.ply``` is very noisy, so we are additionally define the region of interest using MeshLab and upload it to the current folder as ```point_cloud_cropped.ply```. #### Step 3. Transform cropped scene to lie in a unit sphere volume. @@ -85,17 +82,21 @@ Note, now NeuralHaircut supports only the square images. #### Step 5. Obtain hair, silhouette masks and orientation and confidence maps. +For the hair and silhouette masks, the following pretrained models are needed: +- `modnet_photographic_portrait_matting.ckpt`: Download from this [Google Drive](https://drive.google.com/drive/folders/1umYmlCulvIFNaqPjwod1SayFmSRHziyR) and put it in `./MODNet/pretrained` +- `LIP_epoch_149.pth`: Download from this [Google Drive](https://drive.google.com/drive/folders/1E9GutnsqFzF16bC5_DmoSXFIHYuU547L?usp=sharing) and put it in `./CDGNet/snapshots` +Then run: ```bash -python preprocess_custom_data/calc_masks.py --scene_path ./implicit-hair-data/data/SCENE_TYPE/CASE/ --MODNET_ckpt path_to_modnet --CDGNET_ckpt path_to_cdgnet +python preprocess_custom_data/calc_masks.py --scene_path ./implicit-hair-data/data/SCENE_TYPE/CASE/ ``` +After this step in```./implicit-hair-data/data/SCENE_TYPE/CASE``` you would obtain ```hair_mask``` and ```mask```. - +For the orientation and confidence maps, run: ```bash python preprocess_custom_data/calc_orientation_maps.py --img_path ./implicit-hair-data/data/SCENE_TYPE/CASE/image/ --orient_dir ./implicit-hair-data/data/SCENE_TYPE/CASE/orientation_maps --conf_dir ./implicit-hair-data/data/SCENE_TYPE/CASE/confidence_maps ``` - -After this step in```./implicit-hair-data/data/SCENE_TYPE/CASE``` you would obtain ```hair_mask, mask, confidence_maps, orientation_maps```. +After this step in```./implicit-hair-data/data/SCENE_TYPE/CASE``` you would obtain ```confidence_maps``` and ```orientation_maps```. #### Step 6. (Optional) Define views on which you want to train and save it into views.pickle file. diff --git a/preprocess_custom_data/scale_scene_into_sphere.py b/preprocess_custom_data/scale_scene_into_sphere.py index aaa3fd3..fcfc941 100644 --- a/preprocess_custom_data/scale_scene_into_sphere.py +++ b/preprocess_custom_data/scale_scene_into_sphere.py @@ -1,6 +1,7 @@ import trimesh import numpy as np import pickle +import argparse import os diff --git a/preprocess_custom_data/verify_camera.py b/preprocess_custom_data/verify_camera.py new file mode 100644 index 0000000..044e0c1 --- /dev/null +++ b/preprocess_custom_data/verify_camera.py @@ -0,0 +1,60 @@ +import argparse +import numpy as np +import trimesh +import os +import cv2 + +if __name__ == '__main__': + parser = argparse.ArgumentParser() + + parser.add_argument('--cameras_npz', type=str, required=True, help='Path to cameras.npz file') + parser.add_argument('--point_cloud_ply', type=str, required=True, help='Path to point_cloud.ply file') + parser.add_argument('--image_dir', type=str, required=True, help='Path to directory containing images') + parser.add_argument('--output_dir', type=str, required=True, help='Path to output directory') + parser.add_argument('--point_indices', type=str, default=None, help='Comma-separated list of point indices to visualize') + parser.add_argument('--point_radius', type=int, default=4, help='Radius of points in pixels [4]') + parser.add_argument('--seed_offset', type=int, default=0, help='Seed offset for random number generator [0]') + + args = parser.parse_args() + + os.makedirs(args.output_dir, exist_ok=True) + + cameras = np.load(args.cameras_npz)['arr_0'] + + points = np.array(trimesh.load(args.point_cloud_ply).vertices) + + if args.point_indices is not None: + args.point_indices = [int(x) for x in args.point_indices.split(',')] + + for i, image_file in enumerate(sorted(os.listdir(args.image_dir))): + image = cv2.imread(os.path.join(args.image_dir, image_file)) + + width, height = image.shape[1], image.shape[0] + assert(width == height) + + KR = cameras[i, 0:3, 0:3] + Kt = cameras[i, 0:3, 3] + + for j, point in enumerate(points): + if args.point_indices is not None and j not in args.point_indices: + continue + + # Project point into image + uv = KR @ point + Kt + uv /= uv[2] + uv = np.round(uv).astype(int) + + # Choose random color deterministically by index j + np.random.seed(j + args.seed_offset) + color = np.random.randint(0, 255, 3) + + # Rasterize point + r = args.point_radius + for du in range(-r, r+1): + for dv in range(-r, r+1): + if (du * du + dv * dv <= r*r + 3): + if (uv[0] + du >= 0 and uv[0] + du < width and + uv[1] + dv >= 0 and uv[1] + dv < height): + image[uv[1] + dv, uv[0] + du] = color + + cv2.imwrite(os.path.join(args.output_dir, image_file), image) diff --git a/run_strands_optimization.py b/run_strands_optimization.py index 1252a24..e5be0e9 100644 --- a/run_strands_optimization.py +++ b/run_strands_optimization.py @@ -33,7 +33,7 @@ warnings.filterwarnings("ignore") class Runner: - def __init__(self, conf_path, case='CASE_NAME', scene_type='DATASET_TYPE', checkpoint_name=None, hair_conf_path=None, exp_name=None): + def __init__(self, conf_path, case='CASE_NAME', scene_type='DATASET_TYPE', is_continue=False, checkpoint_name=None, hair_conf_path=None, exp_name=None): self.device = torch.device('cuda') @@ -60,7 +60,14 @@ def __init__(self, conf_path, case='CASE_NAME', scene_type='DATASET_TYPE', check if exp_name is not None: date, time = str(datetime.today()).split('.')[0].split(' ') exps_dir = Path('./exps_second_stage') / exp_name / case / Path(conf_path).stem - cur_dir = date + '_' + time + if is_continue: + prev_exps = sorted(exps_dir.iterdir()) + if len(prev_exps) > 0: + cur_dir = prev_exps[-1].name + else: + raise FileNotFoundError(errno.ENOENT, "No previous experiment in directory", exps_dir) + else: + cur_dir = date + '_' + time self.base_exp_dir = exps_dir / cur_dir else: self.base_exp_dir = self.conf['general']['base_exp_dir'] @@ -78,6 +85,7 @@ def __init__(self, conf_path, case='CASE_NAME', scene_type='DATASET_TYPE', check self.iter_step = 0 + self.is_continue = is_continue self.writer = None set_seed(42) @@ -95,7 +103,26 @@ def __init__(self, conf_path, case='CASE_NAME', scene_type='DATASET_TYPE', check if train_conf['pretrain_strands_path']: print('Upload strands!') self.hair_primitives_trainer.load_weights(train_conf['pretrain_hair_path']) - + + # Load checkpoint + latest_model_name = None + if is_continue: + model_list_raw = os.listdir(os.path.join(self.base_exp_dir, 'hair_primitives')) + model_list = [] + for model_name in model_list_raw: + if model_name[-3:] == 'pth' and int(model_name[5:-4]) <= self.end_iter: + model_list.append(model_name) + model_list.sort() + if checkpoint_name is not None and checkpoint_name + ".pth" in model_list: + latest_model_name = checkpoint_name + ".pth" + else: + latest_model_name = model_list[-1] + + if latest_model_name is not None: + logging.info('Find checkpoint: {}'.format(latest_model_name)) + self.hair_primitives_trainer.load_weights(os.path.join(self.base_exp_dir, 'hair_primitives', latest_model_name)) + self.iter_step = int(latest_model_name[5:-4]) + # Backup codes and configs for debug self.file_backup() @@ -191,12 +218,13 @@ def file_backup(self): parser.add_argument('--case', type=str, default='') parser.add_argument('--scene_type', type=str, default='') parser.add_argument('--hair_conf', type=str, default=None, help='Use hair primitives config') + parser.add_argument('--is_continue', action='store_true') parser.add_argument('--checkpoint', type=str, default=None, help='Checkpoint to continue training') parser.add_argument('--exp_name', type=str, default=None) args = parser.parse_args() torch.cuda.set_device(args.gpu) - runner = Runner(args.conf, args.case, args.scene_type, hair_conf_path=args.hair_conf, checkpoint_name=args.checkpoint, exp_name=args.exp_name) + runner = Runner(args.conf, args.case, args.scene_type, hair_conf_path=args.hair_conf, is_continue=args.is_continue, checkpoint_name=args.checkpoint, exp_name=args.exp_name) runner.train() diff --git a/src/multiview_optimization/PIXIE/data/PUT_SMPLX_DATA_HERE b/src/multiview_optimization/PIXIE/data/PUT_SMPLX_DATA_HERE new file mode 100644 index 0000000..e69de29 diff --git a/src/multiview_optimization/concat_pixie.py b/src/multiview_optimization/concat_pixie.py new file mode 100644 index 0000000..d88bcc7 --- /dev/null +++ b/src/multiview_optimization/concat_pixie.py @@ -0,0 +1,35 @@ +import argparse +import os +import pickle +import numpy +import torch + +if __name__ == '__main__': + parser = argparse.ArgumentParser() + + parser.add_argument('--case', default='person_1', type=str) + parser.add_argument('--scene_type', default='monocular', type=str) + parser.add_argument('--path_to_data', default='../../implicit-hair-data/data/', type=str) + + args = parser.parse_args() + + data_dir = os.path.join(args.path_to_data, args.scene_type, args.case) + + data_param = [] + for img_file in os.listdir(os.path.join(data_dir, 'image')): + img_file_stem = img_file.split('.')[0] + # Read individual pixie parameter file + with open(os.path.join(data_dir, 'pixie', img_file_stem,f'{img_file_stem}_param.pkl'), 'rb') as f: + kv = pickle.load(f) + # Convert numpy.ndarray to torch.tensor + for key in kv.keys(): + if type(kv[key]) is numpy.ndarray: + kv[key] = torch.from_numpy(kv[key]) + # Insert a new axis to the front + kv[key] = kv[key].unsqueeze(0) + data_param.append(kv) + + # Write to initialization_pixie + with open(os.path.join(data_dir, 'initialization_pixie'), 'wb') as f: + for p in data_param: + pickle.dump(p, f) diff --git a/src/multiview_optimization/dataset/dataset.py b/src/multiview_optimization/dataset/dataset.py index 61936f4..26e5247 100644 --- a/src/multiview_optimization/dataset/dataset.py +++ b/src/multiview_optimization/dataset/dataset.py @@ -96,7 +96,7 @@ def __init__(self, image_path='', scale_path='', camera_path='' , openpose_kp_pa self.camera_model = None - if fitted_camera_path: + if fitted_camera_path and os.path.exists(fitted_camera_path): self.camera_model = OptimizableCameras(len(imgs_list), pretrain_path=fitted_camera_path) with torch.no_grad(): intrinsics_all, pose_all = self.camera_model(torch.arange(len(imgs_list)), intrinsics_all, pose_all) diff --git a/src/multiview_optimization/readme.md b/src/multiview_optimization/readme.md index c2e0bba..e0dee53 100644 --- a/src/multiview_optimization/readme.md +++ b/src/multiview_optimization/readme.md @@ -2,12 +2,32 @@ Download pre-trained models and data from [SMPLX](https://smpl-x.is.tue.mpg.de/) For more information please follow [PIXIE installation](https://github.com/yfeng95/PIXIE/blob/master/Doc/docs/getting_started.md). -For multiview optimization you need to have the following files ```SMPL-X__FLAME_vertex_ids.npy, smplx_extra_joints.yaml, SMPLX_NEUTRAL_2020.npz``` and change a path to them in ```./utils/config.py``` +For multiview optimization you need to have the following files: +- `SMPL-X__FLAME_vertex_ids.npy`: Download *"MANO and FLAME vertex indices"* from [SMPLX download page](https://smpl-x.is.tue.mpg.de/download.php) +- `smplx_extra_joints.yaml`: Download *"SHAPY Model"* from [SHAPY download page](https://shapy.is.tue.mpg.de/download.php) +- `SMPLX_NEUTRAL_2020.npz`: Download *"SMPL-X 2020"* from [SMPLX download page](https://smpl-x.is.tue.mpg.de/download.php) +and put these files into `./PIXIE/data`. -Note, that you need to obtain [PIXIE initialization](https://github.com/yfeng95/PIXIE) for shape, pose parameters and save it as a dict in ```initialization_pixie``` file (see the structure in [example scene](../../example) for convenience). +Note, that you need to obtain [PIXIE initialization](https://github.com/yfeng95/PIXIE) for shape, pose parameters and save it as a dict in ```initialization_pixie```. +Specifically, assuming that the case name is `person_1` and the image data is stored in `NeuralHaircut/implicit-hair-data/data/monocular/person_1`, run PIXIE on the images: -Furthermore, use [OpenPose](https://github.com/CMU-Perceptual-Computing-Lab/openpose) to obtain 3d keypoints or use only [FaceAlignment](https://github.com/1adrianb/face-alignment) loss in optimization process. +```bash +cd +python demos/demo_fit_face.py -i /implicit-hair-data/data/monocular/person_1/image -s /implicit-hair-data/data/monocular/person_1/pixie --saveParam True +``` +Then, run: +``` +python ./concat_pixie.py --case person_1 --scene_type monocular +``` +which will produce `NeuralHaircut/implicit-hair-data/data/monocular/person_1/initialization_pixie`. + + +Furthermore, use [OpenPose](https://github.com/CMU-Perceptual-Computing-Lab/openpose) to obtain 3d keypoints: +```bash +cd +./build/examples/openpose/openpose.bin -display 0 -render_pose 0 -face -hand -image_dir /implicit-hair-data/data/monocular/person_1/image -write_json /implicit-hair-data/data/monocular/person_1/openpose_kp +``` To obtain FLAME prior run: @@ -21,4 +41,4 @@ To visualize the training process: tensorboard --logdir ./experiments/EXP_NAME ``` -After training put obtained FLAME prior mesh.obj into the dataset folder ```./implicit-hair-data/data/SCENE_TYPE/CASE/head_prior.obj```. \ No newline at end of file +After training put obtained FLAME prior mesh (i.e., the last `*.obj` in `./experiments/fit_person_1_bs_20_train_rot_shape/mesh`) into the dataset folder ```./implicit-hair-data/data/monocular/person_1/head_prior.obj```. diff --git a/src/multiview_optimization/utils/config.py b/src/multiview_optimization/utils/config.py index 431494a..fab554d 100644 --- a/src/multiview_optimization/utils/config.py +++ b/src/multiview_optimization/utils/config.py @@ -9,7 +9,7 @@ cfg = CN() abs_pixie_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..')) -cfg.pixie_dir = '/Vol0/user/v.sklyarova/PIXIE/PIXIE' +cfg.pixie_dir = './PIXIE' cfg.device = 'cuda' cfg.device_id = '0' cfg.pretrained_modelpath = os.path.join(cfg.pixie_dir, 'data', 'pixie_model.tar')