132 example data and scripts not unified

happypose/pose_estimators/cosypose/cosypose/datasets/bop_object_datasets.py

@@ -1,19 +1,22 @@
 import json
 from pathlib import Path
+from typing import Union
 
 
 class BOPObjectDataset:
-    def __init__(self, ds_dir):
+    def __init__(self, ds_dir, label_format: Union[None, str] = None):
         ds_dir = Path(ds_dir)
         infos_file = ds_dir / "models_info.json"
         infos = json.loads(infos_file.read_text())
         objects = []
         for obj_id, bop_info in infos.items():
             obj_id = int(obj_id)
-            obj_label = f"obj_{obj_id:06d}"
-            mesh_path = (ds_dir / obj_label).with_suffix(".ply").as_posix()
+            label = f"obj_{obj_id:06d}"
+            mesh_path = (ds_dir / label).with_suffix(".ply").as_posix()
+            if label_format is not None:
+                label = label_format.format(label=label)
             obj = {
-                "label": obj_label,
+                "label": label,
                 "category": None,
                 "mesh_path": mesh_path,
                 "mesh_units": "mm",

happypose/pose_estimators/cosypose/cosypose/lib3d/rigid_mesh_database.py

@@ -92,7 +92,11 @@
         return {label: obj["n_sym"] for label, obj in self.infos.items()}
 
     def select(self, labels):
-        ids = [self.label_to_id[label] for label in labels]
+        try:
+            ids = [self.label_to_id[label] for label in labels]
+        except KeyError as e:
+            print(e)
+            print("self.label_to_id.keys(): ", list(self.label_to_id.keys()))
         return Meshes(
             infos=[self.infos[label] for label in labels],
             labels=self.labels[ids],

happypose/pose_estimators/cosypose/cosypose/simulator/caching.py

@@ -15,8 +15,10 @@
         self.away_transform = (0, 0, 1000), (0, 0, 0, 1)
 
     def _load_body(self, label):
-        ds_idx = np.where(self.urdf_ds.index["label"] == label)[0].item()
-        object_infos = self.urdf_ds[ds_idx].to_dict()
+        ds_idx = np.where(self.urdf_ds.index["label"] == label)[0]
+        if len(ds_idx) == 0:
+            raise ValueError(f'Label {label} not in {self.urdf_ds.index["label"]}')
+        object_infos = self.urdf_ds[ds_idx.item()].to_dict()
         body = Body.load(
             object_infos["urdf_path"],
             scale=object_infos["scale"],

happypose/pose_estimators/cosypose/cosypose/visualization/multiview.py

@@ -152,7 +152,7 @@
             {"K": K, "TWC": TWC, "resolution": (w, h)},
         )
     renders = renderer.render_scene(list_objects, list_cameras)
-    images = np.stack([render["rgb"] for render in renders])
+    images = np.stack([render.rgb for render in renders])
     if gui:
         time.sleep(100)
     renderer.disconnect()
@@ -201,7 +201,7 @@
             "TWO": preds_with_colors.poses[n].cpu().numpy(),
         }
         list_objects.append(obj)
-    rgb_rendered = renderer.render_scene(list_objects, [camera])[0]["rgb"]
+    rgb_rendered = renderer.render_scene(list_objects, [camera])[0].rgb
     return rgb_rendered
 
 
@@ -212,7 +212,7 @@
         obj["color"] = colormap_rgb[obj["label"]]
         obj["TWO"] = np.linalg.inv(TWC) @ obj["TWO"]
     camera["TWC"] = np.eye(4)
-    rgb_rendered = renderer.render_scene(objects, [camera])[0]["rgb"]
+    rgb_rendered = renderer.render_scene(objects, [camera])[0].rgb
     return rgb_rendered
 
 

happypose/toolbox/datasets/datasets_cfg.py

@@ -249,7 +249,12 @@
 def make_object_dataset(ds_name: str) -> RigidObjectDataset:
     # BOP original models
 
-    if ds_name == "tless.cad":
+    if ds_name == "tless":
+        ds: RigidObjectDataset = BOPObjectDataset(
+            BOP_DS_DIR / "tless/models_cad",
+            label_format="tless-{label}",
+        )
+    elif ds_name == "tless.cad":
         ds: RigidObjectDataset = BOPObjectDataset(
             BOP_DS_DIR / "tless/models_cad",
             label_format="tless-{label}",

happypose/toolbox/renderer/bullet_scene_renderer.py

@@ -7,6 +7,7 @@
 
 # from happypose.toolbox.datasets.datasets_cfg import UrdfDataset
 from happypose.pose_estimators.cosypose.cosypose.datasets.urdf_dataset import (
+    OneUrdfDataset,
     UrdfDataset,
 )
 
@@ -31,7 +32,7 @@ def __init__(
         gpu_renderer=True,
         gui=False,
     ):
-        if isinstance(asset_dataset, UrdfDataset):
+        if isinstance(asset_dataset, (UrdfDataset, OneUrdfDataset)):
             self.urdf_ds = asset_dataset
         elif isinstance(asset_dataset, RigidObjectDataset):
             # Build urdfs files from RigidObjectDataset
@@ -43,7 +44,7 @@ def __init__(
             self.urdf_ds.index["scale"] = asset_dataset[0].scale
         else:
             raise TypeError(
-                f"asset_dataset of type {type(asset_dataset)} should be either UrdfDataset or RigidObjectDataset"
+                f"asset_dataset of type {type(asset_dataset)} should be either OneUrdfDataset/UrdfDataset or RigidObjectDataset"
             )
         self.connect(gpu_renderer=gpu_renderer, gui=gui)
         self.body_cache = BodyCache(self.urdf_ds, self.client_id)

happypose/toolbox/renderer/types.py

@@ -160,6 +160,10 @@
     scale: float = 1
     positioning_function: Optional[NodeFunction] = None
 
+    def __post_init__(self):
+        if not isinstance(self.TWO, Transform):
+            self.TWO = Transform(self.TWO)
+
     def set_node_material_and_transparency(
         self,
         node_path: p3d.core.NodePath,