first commit

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.

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
 }

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
 }

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

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

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)

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