fix small bugs

preprocess_custom_data/copy_checkpoints.py

@@ -1,14 +1,15 @@
 from shutil import copyfile
 import os
 from pathlib import Path
+import argparse
 
 def main(args):
 
     exp_name = args.exp_name
     case = args.case
     conf_path = args.conf_path
 
-    exps_dir = Path('./exps_first_stage') / exp_name / case / Path(conf_path).stem
+    exps_dir = Path(args.exp_path) / exp_name / case / Path(conf_path).stem
     prev_exps = sorted(exps_dir.iterdir())
     cur_dir = prev_exps[-1].name      
 
@@ -19,19 +20,28 @@ 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]
 
-    print(f'Copy obtained from the first stage checkpoint, hair and head geometry to folder ./implicit-hair-data/data/{case}')
+    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]
+
+    path_to_data = args.path_to_data
+    scene_type = args.scene_type
+
+
+    print(f'Copy obtained from the first stage checkpoint, hair and head geometry to folder {path_to_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')
+
+    print(path_to_data, os.path.join(path_to_mesh, last_hair), f'{path_to_data}/{scene_type}/{case}/final_hair.ply')
+
+
+    copyfile(os.path.join(path_to_mesh, last_hair), f'{path_to_data}/{scene_type}/{case}/final_hair.ply')
+    copyfile(os.path.join(path_to_mesh, last_head), f'{path_to_data}/{scene_type}/{case}/final_head.ply')
+    copyfile(os.path.join(path_to_ckpt, last_ckpt), f'{path_to_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 {path_to_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'{path_to_data}/{scene_type}/{case}/fitted_cameras.pth')
 
 
 
@@ -43,6 +53,12 @@ def main(args):
     parser.add_argument('--case', default='person_1', type=str)
 
     parser.add_argument('--exp_name', default='first_stage_reconctruction_person_1', type=str)  
+
+    parser.add_argument('--exp_path', default='./exps_first_stage', type=str) 
+
+    parser.add_argument('--scene_type', default='monocular', type=str, choices=['h3ds', 'monocular']) 
+
+    parser.add_argument('--path_to_data', default='./implicit-hair-data/data/', type=str)  
 
 
     args, _ = parser.parse_known_args()

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

@@ -13,7 +13,7 @@
 from pytorch3d.io import load_ply, save_ply
 
 sys.path.append(os.path.join(sys.path[0], '..'))
-from src.models.dataset import Dataset, MonocularDataset
+from src.models.dataset import Dataset, MonocularDataset, CustomDataset
 
 import argparse
 import yaml
@@ -63,7 +63,7 @@ def main(args):
 
     scene_type = args.scene_type
     case = args.case
-    save_path = f'./implicit-hair-data/data/{scene_type}/{case}'
+    save_path = os.path.join(args.path_to_data, f'{scene_type}/{case}')
 
     # upload mesh hair and bust
     verts, faces = load_ply(os.path.join(save_path, 'final_hair.ply'))
@@ -73,8 +73,10 @@ def main(args):
 
     mesh_head =  Meshes(verts=[(verts_).float().to(args.device)], faces=[faces_.to(args.device)])
 
+    H, W = args.img_size
+
     raster_settings_mesh = RasterizationSettings(
-                        image_size=args.img_size, 
+                        image_size=(int(H), int(W)), 
                         blur_radius=0.000, 
                         faces_per_pixel=1, 
                     )
@@ -83,7 +85,7 @@ def main(args):
     R = torch.ones(1, 3, 3)
     t = torch.ones(1, 3)
     cam_intr = torch.ones(1, 4, 4)
-    size = torch.tensor([args.img_size, args.img_size ]).to(args.device)
+    size = torch.tensor([int(H), int(W)]).to(args.device)
 
     cam = cameras_from_opencv_projection(
                                         camera_matrix=cam_intr.cuda(), 
@@ -110,8 +112,13 @@ def main(args):
 
     if scene_type == 'h3ds':
         dataset = Dataset(conf['dataset'])
+
+    elif scene_type=='monocular':    
+        dataset = CustomDataset(conf['dataset'])
+
     else:
-        dataset = MonocularDataset(conf['dataset'])
+        pass
+
 
 
     # add upper cameras for h3ds data as they haven't got such views
@@ -145,14 +152,14 @@ def main(args):
     parser = argparse.ArgumentParser(conflict_handler='resolve')
 
     parser.add_argument('--conf_path', default='./configs/monocular/neural_strands.yaml', type=str)
-
+    parser.add_argument('--path_to_data', default='./implicit_hair_data/data', type=str)
     parser.add_argument('--case', default='person_1', type=str)
 
     parser.add_argument('--scene_type', default='monocular', type=str) 
 
     parser.add_argument('--device', default='cuda', type=str)
 
-    parser.add_argument('--img_size', default=2160, type=int)
+    parser.add_argument('--img_size', default=[2160, 2160], nargs="+")
     parser.add_argument('--n_views', default=2, type=int)
 
     args, _ = parser.parse_known_args()

preprocess_custom_data/scale_scene_into_sphere.py

@@ -3,6 +3,7 @@
 import pickle
 
 import os
+import argparse
 
 
 

src/utils/util.py

@@ -6,11 +6,11 @@
 import random
 
 
-def scale_mat(mat, scale_factor):
-    mat[0, 0] /= scale_factor
-    mat[1, 1] /= scale_factor
-    mat[0, 2] /= scale_factor
-    mat[1, 2] /= scale_factor
+def scale_mat(mat, scale_factor_H, scale_factor_W):
+    mat[0, 0] /= scale_factor_W
+    mat[1, 1] /= scale_factor_H
+    mat[0, 2] /= scale_factor_W
+    mat[1, 2] /= scale_factor_H
     return mat
 
 def positional_encoding(