Merge pull request #27 from prithvirb/hot3d

[cnos] add cad model template rendering for hot3d

src/poses/pyrender.py

@@ -71,6 +71,7 @@ def render(
     parser.add_argument("radius", nargs="?", type=float, default=1, help="Distance from camera to object")
     args = parser.parse_args()
     print(args)
+    is_hot3d = "hot3d" in args.output_dir
     os.environ["CUDA_VISIBLE_DEVICES"] = args.gpus_devices
     poses = np.load(args.obj_pose)
     # we can increase high energy for lightning but it's simpler to change just scale of the object to meter
@@ -93,15 +94,22 @@ def render(
 
     # load mesh to meter
     mesh = trimesh.load_mesh(args.cad_path)
+
     diameter = get_obj_diameter(mesh)
-    if diameter > 100: # object is in mm
+    if not is_hot3d and diameter > 100: # object is in mm
         mesh.apply_scale(0.001)
+
     if is_tless:
         # setting uniform colors for mesh
         color = 0.4
         mesh.visual.face_colors = np.ones((len(mesh.faces), 3)) * color
         mesh.visual.vertex_colors = np.ones((len(mesh.vertices), 3)) * color
         mesh = pyrender.Mesh.from_trimesh(mesh, smooth=False)
+    elif is_hot3d:
+        mesh = pyrender.Mesh.from_trimesh(list(mesh.geometry.values())[0])
+        ## Rescale mesh vertices to be as same unit as the other dataset values
+        geometry = mesh.primitives[0].positions
+        mesh.primitives[0].positions = geometry * 0.001
     else:
         mesh = pyrender.Mesh.from_trimesh(as_mesh(mesh))
     os.makedirs(args.output_dir, exist_ok=True)

src/scripts/render_template_with_pyrender.py

@@ -35,6 +35,7 @@ def call_pyrender(
 
     # command = f"blenderproc run ./src/poses/blenderproc.py {cad_path} {obj_pose_path} {output_dir} {gpus_devices}"
     command = f"python -m src.poses.pyrender {cad_path} {obj_pose_path} {output_dir} {gpus_devices}"
+    print(f"executing command:\n{command}")
     if disable_output:
         command += " true"
     else:
@@ -75,11 +76,19 @@ def render(cfg: DictConfig) -> None:
         obj_pose_path = f"{dataset_save_dir}/template_poses.npy"
         np.save(obj_pose_path, template_poses)
 
+        cad_file_ext = None
         if dataset_name in ["tless"]:
+            cad_file_ext = ".ply"
             cad_dir = os.path.join(cfg.data.root_dir, dataset_name, "models/models_cad")
             if not os.path.exists(cad_dir):
                 cad_dir = os.path.join(cfg.data.root_dir, dataset_name, "models_cad")
+        elif dataset_name == "hot3d":
+            cad_file_ext = ".glb"
+            cad_dir = os.path.join(cfg.data.root_dir, dataset_name, "object_models")
+            if not os.path.exists(cad_dir):
+                raise Exception(f"cad folder not found. cad_dir: {cad_dir}")            
         else:
+            cad_file_ext = ".ply"
             cad_dir = os.path.join(cfg.data.root_dir, dataset_name, "models/models")
             if not os.path.exists(cad_dir):
                 cad_dir = os.path.join(cfg.data.root_dir, dataset_name, "models")
@@ -89,14 +98,16 @@ def render(cfg: DictConfig) -> None:
             [
                 int(name[4:][:-4])
                 for name in os.listdir(cad_dir)
-                if name.endswith(".ply")
+                if name.endswith(cad_file_ext)
             ]
         )
+        print(f"num object_ids found: {len(object_ids)}")
+        print(f"object_ids: {object_ids}")
         for object_id in object_ids:
             cad_paths.append(
                 os.path.join(
                     cad_dir,
-                    "obj_{:06d}.ply".format(object_id),
+                    f"obj_{object_id:06d}{cad_file_ext}",
                 )
             )
             output_dirs.append(