Fixes and enhancements

.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/

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

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()

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()

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()

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')
 
 
 

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):
 

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)

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.

preprocess_custom_data/scale_scene_into_sphere.py

@@ -1,6 +1,7 @@
 import trimesh
 import numpy as np
 import pickle
+import argparse
 
 import os