Spaces:
Runtime error
Runtime error
File size: 10,127 Bytes
e2ebf5f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 |
import argparse
from sugar.sugar_utils.general_utils import str2bool
from sugar.sugar_trainers.coarse_density import coarse_training_with_density_regularization
from sugar.sugar_trainers.coarse_sdf import coarse_training_with_sdf_regularization
from sugar.sugar_extractors.coarse_mesh import extract_mesh_from_coarse_sugar
from sugar.sugar_trainers.refine import refined_training
from sugar.sugar_extractors.refined_mesh import extract_mesh_and_texture_from_refined_sugar
class AttrDict(dict):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self.__dict__ = self
def main(scene_path="data/custom/", checkpoint_path="output/custom/", regularization_type='density', export_uv_textured_mesh=True, export_ply=True, low_poly=True, refinement_time='short'):
# ----- Parser -----
parser = argparse.ArgumentParser(description='Script to optimize a full SuGaR model.')
# Data and vanilla 3DGS checkpoint
parser.add_argument('-s', '--scene_path',
type=str,
help='(Required) path to the scene data to use.')
parser.add_argument('-c', '--checkpoint_path',
type=str,
help='(Required) path to the vanilla 3D Gaussian Splatting Checkpoint to load.')
parser.add_argument('-i', '--iteration_to_load',
type=int, default=7000,
help='iteration to load.')
# Regularization for coarse SuGaR
parser.add_argument('-r', '--regularization_type', type=str,
help='(Required) Type of regularization to use for coarse SuGaR. Can be "sdf" or "density". '
'For reconstructing detailed objects centered in the scene with 360° coverage, "density" provides a better foreground mesh. '
'For a stronger regularization and a better balance between foreground and background, choose "sdf".')
# Extract mesh
parser.add_argument('-l', '--surface_level', type=float, default=0.3,
help='Surface level to extract the mesh at. Default is 0.3')
parser.add_argument('-v', '--n_vertices_in_mesh', type=int, default=1_000_000,
help='Number of vertices in the extracted mesh.')
parser.add_argument('-b', '--bboxmin', type=str, default=None,
help='Min coordinates to use for foreground.')
parser.add_argument('-B', '--bboxmax', type=str, default=None,
help='Max coordinates to use for foreground.')
parser.add_argument('--center_bbox', type=str2bool, default=True,
help='If True, center the bbox. Default is False.')
# Parameters for refined SuGaR
parser.add_argument('-g', '--gaussians_per_triangle', type=int, default=1,
help='Number of gaussians per triangle.')
parser.add_argument('-f', '--refinement_iterations', type=int, default=15_000,
help='Number of refinement iterations.')
# (Optional) Parameters for textured mesh extraction
parser.add_argument('-t', '--export_uv_textured_mesh', type=str2bool, default=True,
help='If True, will export a textured mesh as an .obj file from the refined SuGaR model. '
'Computing a traditional colored UV texture should take less than 10 minutes.')
parser.add_argument('--square_size',
default=10, type=int, help='Size of the square to use for the UV texture.')
parser.add_argument('--postprocess_mesh', type=str2bool, default=False,
help='If True, postprocess the mesh by removing border triangles with low-density. '
'This step takes a few minutes and is not needed in general, as it can also be risky. '
'However, it increases the quality of the mesh in some cases, especially when an object is visible only from one side.')
parser.add_argument('--postprocess_density_threshold', type=float, default=0.1,
help='Threshold to use for postprocessing the mesh.')
parser.add_argument('--postprocess_iterations', type=int, default=5,
help='Number of iterations to use for postprocessing the mesh.')
# (Optional) PLY file export
parser.add_argument('--export_ply', type=str2bool, default=True,
help='If True, export a ply file with the refined 3D Gaussians at the end of the training. '
'This file can be large (+/- 500MB), but is needed for using the dedicated viewer. Default is True.')
# (Optional) Default configurations
parser.add_argument('--low_poly', type=str2bool, default=False,
help='Use standard config for a low poly mesh, with 200k vertices and 6 Gaussians per triangle.')
parser.add_argument('--high_poly', type=str2bool, default=False,
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'.")
# Evaluation split
parser.add_argument('--eval', type=str2bool, default=True, help='Use eval split.')
# GPU
parser.add_argument('--gpu', type=int, default=0, help='Index of GPU device to use.')
# Parse arguments
args = parser.parse_args()
if args.low_poly:
args.n_vertices_in_mesh = 200_000
args.gaussians_per_triangle = 6
print('Using low poly config.')
if args.high_poly:
args.n_vertices_in_mesh = 1_000_000
args.gaussians_per_triangle = 1
print('Using high poly config.')
if args.refinement_time == 'short':
args.refinement_iterations = 2_000
print('Using short refinement time.')
if args.refinement_time == 'medium':
args.refinement_iterations = 7_000
print('Using medium refinement time.')
if args.refinement_time == 'long':
args.refinement_iterations = 15_000
print('Using long refinement time.')
if args.export_uv_textured_mesh:
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.')
if scene_path:
args.scene_path = scene_path
if checkpoint_path:
args.checkpoint_path = checkpoint_path
if regularization_type:
args.regularization_type = regularization_type
if export_uv_textured_mesh:
args.export_uv_textured_mesh = export_uv_textured_mesh
if export_ply:
args.export_ply = export_ply
if low_poly:
args.low_poly = low_poly
if refinement_time:
args.refinement_time = refinement_time
# ----- 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,
'eval': args.eval,
'estimation_factor': 0.2,
'normal_factor': 0.2,
'gpu': args.gpu,
})
if args.regularization_type == 'sdf':
coarse_sugar_path = coarse_training_with_sdf_regularization(coarse_args)
elif args.regularization_type == 'density':
coarse_sugar_path = coarse_training_with_density_regularization(coarse_args)
else:
raise ValueError(f'Unknown regularization type: {args.regularization_type}')
# ----- Extract mesh from coarse SuGaR -----
coarse_mesh_args = AttrDict({
'scene_path': args.scene_path,
'checkpoint_path': args.checkpoint_path,
'iteration_to_load': args.iteration_to_load,
'coarse_model_path': coarse_sugar_path,
'surface_level': args.surface_level,
'decimation_target': args.n_vertices_in_mesh,
'mesh_output_dir': None,
'bboxmin': args.bboxmin,
'bboxmax': args.bboxmax,
'center_bbox': args.center_bbox,
'gpu': args.gpu,
'eval': args.eval,
'use_centers_to_extract_mesh': False,
'use_marching_cubes': False,
'use_vanilla_3dgs': False,
})
coarse_mesh_path = extract_mesh_from_coarse_sugar(coarse_mesh_args)[0]
# ----- Refine SuGaR -----
refined_args = AttrDict({
'scene_path': args.scene_path,
'checkpoint_path': args.checkpoint_path,
'mesh_path': coarse_mesh_path,
'output_dir': None,
'iteration_to_load': args.iteration_to_load,
'normal_consistency_factor': 0.1,
'gaussians_per_triangle': args.gaussians_per_triangle,
'n_vertices_in_fg': args.n_vertices_in_mesh,
'refinement_iterations': args.refinement_iterations,
'bboxmin': args.bboxmin,
'bboxmax': args.bboxmax,
'export_ply': args.export_ply,
'eval': args.eval,
'gpu': args.gpu,
})
refined_sugar_path = refined_training(refined_args)
# ----- Extract mesh and texture from refined SuGaR -----
if args.export_uv_textured_mesh:
refined_mesh_args = AttrDict({
'scene_path': args.scene_path,
'iteration_to_load': args.iteration_to_load,
'checkpoint_path': args.checkpoint_path,
'refined_model_path': refined_sugar_path,
'mesh_output_dir': None,
'n_gaussians_per_surface_triangle': args.gaussians_per_triangle,
'square_size': args.square_size,
'eval': args.eval,
'gpu': args.gpu,
'postprocess_mesh': args.postprocess_mesh,
'postprocess_density_threshold': args.postprocess_density_threshold,
'postprocess_iterations': args.postprocess_iterations,
})
refined_mesh_path = extract_mesh_and_texture_from_refined_sugar(refined_mesh_args)
if __name__ == "__main__":
main()
|