diff --git a/sugar_extractors/refined_mesh.py b/sugar_extractors/refined_mesh.py index eb4a815..d0df970 100644 --- a/sugar_extractors/refined_mesh.py +++ b/sugar_extractors/refined_mesh.py @@ -47,8 +47,13 @@ def extract_mesh_and_texture_from_refined_sugar(args): mesh_save_path = os.path.join(mesh_output_dir, mesh_save_path) scene_name = source_path.split('/')[-2] if len(source_path.split('/')[-1]) == 0 else source_path.split('/')[-1] - sugar_mesh_path = os.path.join('./output/coarse_mesh/', scene_name, + + # --- Coarse mesh parameters --- + if args.coarse_mesh_path is None: + sugar_mesh_path = os.path.join('./output/coarse_mesh/', scene_name, refined_model_path.split('/')[-2].split('_normalconsistency')[0].replace('sugarfine', 'sugarmesh') + '.ply') + else: + sugar_mesh_path = args.coarse_mesh_path if args.square_size is None: if n_gaussians_per_surface_triangle == 1: diff --git a/train.py b/train.py index c7a4309..8a33ff0 100644 --- a/train.py +++ b/train.py @@ -1,4 +1,5 @@ import argparse +from pathlib import Path from sugar_utils.general_utils import str2bool from sugar_trainers.coarse_density import coarse_training_with_density_regularization from sugar_trainers.coarse_sdf import coarse_training_with_sdf_regularization @@ -79,7 +80,9 @@ def __init__(self, *args, **kwargs): help='Use standard config for a high poly mesh, with 1M vertices and 1 Gaussians per triangle.') parser.add_argument('--refinement_time', type=str, default=None, help="Default configs for time to spend on refinement. Can be 'short', 'medium' or 'long'.") - + parser.add_argument('-o', '--output', type=str, default="./output/", + help='Output directory for the meshes') + # Evaluation split parser.add_argument('--eval', type=str2bool, default=True, help='Use eval split.') @@ -109,13 +112,14 @@ def __init__(self, *args, **kwargs): print('Will export a UV-textured mesh as an .obj file.') if args.export_ply: print('Will export a ply file with the refined 3D Gaussians at the end of the training.') + output_dir = Path(args.output) # ----- Optimize coarse SuGaR ----- coarse_args = AttrDict({ 'checkpoint_path': args.checkpoint_path, 'scene_path': args.scene_path, 'iteration_to_load': args.iteration_to_load, - 'output_dir': None, + 'output_dir': output_dir/"coarse", 'eval': args.eval, 'estimation_factor': 0.2, 'normal_factor': 0.2, @@ -137,7 +141,7 @@ def __init__(self, *args, **kwargs): 'coarse_model_path': coarse_sugar_path, 'surface_level': args.surface_level, 'decimation_target': args.n_vertices_in_mesh, - 'mesh_output_dir': None, + 'mesh_output_dir': output_dir/"coarse", 'bboxmin': args.bboxmin, 'bboxmax': args.bboxmax, 'center_bbox': args.center_bbox, @@ -155,7 +159,7 @@ def __init__(self, *args, **kwargs): 'scene_path': args.scene_path, 'checkpoint_path': args.checkpoint_path, 'mesh_path': coarse_mesh_path, - 'output_dir': None, + 'output_dir': output_dir/"refined", 'iteration_to_load': args.iteration_to_load, 'normal_consistency_factor': 0.1, 'gaussians_per_triangle': args.gaussians_per_triangle, @@ -176,8 +180,9 @@ def __init__(self, *args, **kwargs): 'scene_path': args.scene_path, 'iteration_to_load': args.iteration_to_load, 'checkpoint_path': args.checkpoint_path, + 'coarse_mesh_path': coarse_mesh_path, 'refined_model_path': refined_sugar_path, - 'mesh_output_dir': None, + 'mesh_output_dir': output_dir/"refined", 'n_gaussians_per_surface_triangle': args.gaussians_per_triangle, 'square_size': args.square_size, 'eval': args.eval,