diff --git a/README.md b/README.md index bfe14ef..c4fe8d5 100644 --- a/README.md +++ b/README.md @@ -19,11 +19,11 @@ MeshDiffusion is a diffusion model for generating 3D meshes with a direct parame - Pytorch3D -Install https://github.com/NVlabs/nvdiffrec +Follow the instructions to install requirements for nvdiffrec: https://github.com/NVlabs/nvdiffrec ### Pretrained Models -Download the files from +Download the model checkpoints from https://drive.google.com/drive/folders/15IjbUM1tQf8gS0YsRqY5ZbMs-leJgoJ0?usp=sharing. ## Inference @@ -41,10 +41,10 @@ Later run ``` cd nvdiffrec -python eval.py --config $DMTET_CONFIG --sample-path $SAMPLE_PATH +python eval.py --config $DMTET_CONFIG --sample-path $SAMPLE_PATH [--deform-scale $DEFORM_SCALE] ``` -where `$SAMPLE_PATH` is the generated sample `.npy` file in `$OUTPUT_PATH`. +where `$SAMPLE_PATH` is the generated sample `.npy` file in `$OUTPUT_PATH`, and `$DEFORM_SCALE` is the scale of deformation of tet vertices set for the DMTet dataset (we use 3.0 for resolution 64 as default; change the value for your own datasets). ### Single-view Conditional Generation @@ -53,9 +53,11 @@ First fit a DMTet from a single view of a mesh ``` cd nvdiffrec -python fit_singleview.py --mesh-path $MESH_PATH --angle-ind $ANGLE_IND --out-dir $OUT_DIR --validate $VALIDATE +python fit_singleview.py --config $DMTET_CONFIG --mesh-path $MESH_PATH --angle-ind $ANGLE_IND --out-dir $OUT_DIR --validate $VALIDATE ``` +where `$ANGLE_IND` is an integer (0 to 50) controlling the z-axis rotation of the object. Set `$VALIDATE` to 1 if visualization of fitted DMTets is needed. + Then use the trained diffusion model to complete the occluded regions ``` @@ -68,13 +70,15 @@ python main_diffusion.py --mode=cond_gen --config=$DIFFUSION_CONFIG \ --config.eval.batch_size=$EVAL_BATCH_SIZE ``` -Now visualize the completed meshes +Now store the completed meshes as `.obj` files in `$SAMPLE_PATH` ``` cd nvdiffrec -python eval.py --config $DMTET_CONFIG --sample-path $SAMPLE_PATH +python eval.py --config $DMTET_CONFIG --sample-path $SAMPLE_PATH --deform-scale $DEFORM_SCALE ``` +Caution: the deformation scale should be consistent for single view fitting and the diffusion model. Check before you run conditional generation. + ## Training For ShapeNet, first create a list of paths of all ground-truth meshes and store them as a json file under `./nvdiffrec/data/shapenet_json`. @@ -86,9 +90,7 @@ cd nvdiffrec python fit_dmtets.py --config $DMTET_CONFIG --out-dir $DMTET_DATA_PATH ``` - - -Create a meta file for diffusion model training: +Create a meta file of all dmtet grid file locations for diffusion model training: ``` cd ../metadata/ @@ -104,6 +106,8 @@ python main_diffusion.py --mode=train --config=$DIFFUSION_CONFIG \ --config.data.filter_meta_path=$TRAIN_SPLIT_FILE ``` +where `$TRAIN_SPLIT_FILE` is a json list of indices to be included in the training set. Examples in `metadata/train_split/`. + ## Texture Completion Follow the instructions in https://github.com/TEXTurePaper/TEXTurePaper and create text-conditioned textures for the generated meshes. diff --git a/nvdiffrec/configs/res128.json b/nvdiffrec/configs/res128.json index 51bba90..803ab6c 100644 --- a/nvdiffrec/configs/res128.json +++ b/nvdiffrec/configs/res128.json @@ -14,5 +14,6 @@ "background" : "white", "envmap": "./data/irrmaps/aerodynamics_workshop_2k.hdr", "tet_path": "./data/tets/128_tets_cropped.npz", - "cropped": true + "first_stage_deform": 0.45, + "second_stage_deform": 1.5 } \ No newline at end of file diff --git a/nvdiffrec/configs/res64.json b/nvdiffrec/configs/res64.json index 6abca71..68131d7 100644 --- a/nvdiffrec/configs/res64.json +++ b/nvdiffrec/configs/res64.json @@ -14,5 +14,6 @@ "background" : "white", "envmap": "./data/irrmaps/aerodynamics_workshop_2k.hdr", "tet_path": "./data/tets/64_tets_cropped.npz", - "cropped": true + "first_stage_deform": 2.0, + "second_stage_deform": 3.0 } \ No newline at end of file diff --git a/nvdiffrec/eval.py b/nvdiffrec/eval.py index 565ee0c..08fec4c 100644 --- a/nvdiffrec/eval.py +++ b/nvdiffrec/eval.py @@ -376,8 +376,10 @@ if __name__ == "__main__": # Setup geometry for optimization resolution = FLAGS.dmtet_grid - geometry = DMTetGeometry(resolution, FLAGS.mesh_scale, FLAGS) - geometry.deform_scale = FLAGS.deform_scale + geometry = DMTetGeometry( + resolution, FLAGS.mesh_scale, FLAGS, + deform_scale=FLAGS.deform_scale + ) mask = torch.load(f'../data/grid_mask_{resolution}.pt').view(1, resolution, resolution, resolution).to("cuda") diff --git a/nvdiffrec/eval_traj.py b/nvdiffrec/eval_traj.py index f0237a8..f0d7819 100644 --- a/nvdiffrec/eval_traj.py +++ b/nvdiffrec/eval_traj.py @@ -375,8 +375,10 @@ if __name__ == "__main__": # Setup geometry for optimization resolution = FLAGS.dmtet_grid - geometry = DMTetGeometry(resolution, FLAGS.mesh_scale, FLAGS) - geometry.deform_scale = FLAGS.deform_scale + geometry = DMTetGeometry( + resolution, FLAGS.mesh_scale, FLAGS, + deform_scale=FLAGS.deform_scale + ) mask = torch.load(f'../data/grid_mask_{resolution}.pt').view(1, resolution, resolution, resolution).to("cuda") diff --git a/nvdiffrec/fit_dmtets.py b/nvdiffrec/fit_dmtets.py index 1a09025..e798518 100644 --- a/nvdiffrec/fit_dmtets.py +++ b/nvdiffrec/fit_dmtets.py @@ -708,7 +708,10 @@ if __name__ == "__main__": # ============================================================================================== # Setup geometry for optimization - geometry = DMTetGeometry(FLAGS.dmtet_grid, FLAGS.mesh_scale, FLAGS) + geometry = DMTetGeometry( + FLAGS.dmtet_grid, FLAGS.mesh_scale, FLAGS, + deform_scale=FLAGS.first_stage_deform + ) # Setup textures, make initial guess from reference if possible if not FLAGS.normal_only: @@ -759,18 +762,16 @@ if __name__ == "__main__": # # ============================================================================================== # # Pass 2: Finetune deformation with fixed topology # # ============================================================================================== - geometry = DMTetGeometryFixedTopo(geometry, base_mesh, FLAGS.dmtet_grid, FLAGS.mesh_scale, FLAGS) - + geometry = DMTetGeometryFixedTopo( + geometry, base_mesh, FLAGS.dmtet_grid, FLAGS.mesh_scale, FLAGS, + deform_scale=FLAGS.second_stage_deform + ) geometry.sdf_sign.requires_grad = False geometry.sdf_abs.requires_grad = False geometry.deform.requires_grad = True - # geometry.deform.data[:] = geometry.deform * 2.0 / 3.0 - # geometry.deform_scale = 3.0 - - geometry.deform.data[:] = geometry.deform * 0.45 / 1.5 - geometry.deform_scale = 1.5 + geometry.deform.data[:] = geometry.deform * FLAGS.first_stage_deform / FLAGS.second_stage_deform if FLAGS.use_ema: geometry.sdf_sign.data[:] = torch.sign(old_geometry.sdf_ema) diff --git a/nvdiffrec/fit_singleview.py b/nvdiffrec/fit_singleview.py index 2a9753d..8f57a96 100644 --- a/nvdiffrec/fit_singleview.py +++ b/nvdiffrec/fit_singleview.py @@ -705,7 +705,10 @@ if __name__ == "__main__": # ============================================================================================== # Setup geometry for optimization - geometry = DMTetGeometry(FLAGS.dmtet_grid, FLAGS.mesh_scale, FLAGS) + geometry = DMTetGeometry( + FLAGS.dmtet_grid, FLAGS.mesh_scale, FLAGS, + deform_scale=FLAGS.first_stage_deform + ) # Setup textures, make initial guess from reference if possible if not FLAGS.normal_only: @@ -757,15 +760,17 @@ if __name__ == "__main__": base_mesh = xatlas_uvmap(glctx, geometry, mat, FLAGS) - geometry = DMTetGeometryFixedTopo(geometry, base_mesh, FLAGS.dmtet_grid, FLAGS.mesh_scale, FLAGS) + geometry = DMTetGeometryFixedTopo( + geometry, base_mesh, FLAGS.dmtet_grid, FLAGS.mesh_scale, FLAGS, + deform_scale=FLAGS.second_stage_deform + ) geometry.sdf_sign.requires_grad = False geometry.sdf_abs.requires_grad = False geometry.deform.requires_grad = True - geometry.deform.data[:] = geometry.deform * 2.0 / 3.0 - geometry.deform_scale = 3.0 + geometry.deform.data[:] = geometry.deform * FLAGS.first_stage_deform / FLAGS.second_stage_deform if FLAGS.use_ema: geometry.sdf_sign.data[:] = torch.sign(old_geometry.sdf_ema) ### use ema