Merge pull request #142 from agimus-project/132-example-data-and-scri…

…pts-not-unified

132 example data and scripts not unified

docs/book/cosypose/download_data.md

@@ -34,8 +34,8 @@ Notes:
 - The URDF files were obtained using these commands (requires `meshlab` to be installed):
 
   ```sh
-  python -m happypose.pose_estimators.cosypose.cosypose.scripts.convert_models_to_urdf --models=ycbv
-  python -m happypose.pose_estimators.cosypose.cosypose.scripts.convert_models_to_urdf --models=tless.cad
+  python -m happypose.pose_estimators.cosypose.cosypose.scripts.convert_bop_ds_to_urdf --ds_name=ycbv
+  python -m happypose.pose_estimators.cosypose.cosypose.scripts.convert_bop_ds_to_urdf --ds_name=tless.cad
   ```
 
 - Compatibility models were obtained using the following script:

happypose/pose_estimators/cosypose/cosypose/datasets/urdf_dataset.py

@@ -1,18 +1,29 @@
+import json
 from pathlib import Path
 
 import pandas as pd
 
 
 class UrdfDataset:
-    def __init__(self, ds_dir):
-        ds_dir = Path(ds_dir)
+    def __init__(self, urdf_ds_dir, label_filename="objname2label.json"):
+        urdf_ds_dir = Path(urdf_ds_dir)
+        label_path = urdf_ds_dir / label_filename
+        if label_path.exists():
+            with label_path.open() as fp:
+                objname2label = json.load(fp)
+        else:
+            objname2label = None
         index = []
-        for urdf_dir in Path(ds_dir).iterdir():
-            urdf_paths = list(urdf_dir.glob("*.urdf"))
+        for obj_dir in urdf_ds_dir.iterdir():
+            urdf_paths = list(obj_dir.glob("*.urdf"))
             if len(urdf_paths) == 1:
                 urdf_path = urdf_paths[0]
+                if objname2label is None:
+                    label = obj_dir.name
+                else:
+                    label = objname2label[obj_dir.name]
                 infos = {
-                    "label": urdf_dir.name,
+                    "label": label,
                     "urdf_path": urdf_path.as_posix(),
                     "scale": 1.0,
                 }

happypose/pose_estimators/cosypose/cosypose/integrated/pose_estimator.py

@@ -147,7 +147,7 @@ def run_inference_pipeline(
         coarse_estimates: Optional[PoseEstimatesType] = None,
         detection_th: float = 0.7,
         mask_th: float = 0.8,
-        labels_to_keep: List[str] = None,
+        labels_to_keep: Optional[List[str]] = None,
     ) -> Tuple[PoseEstimatesType, dict]:
         timing_str = ""
         timer = SimpleTimer()

happypose/pose_estimators/cosypose/cosypose/libmesh/urdf_utils.py

@@ -1,34 +1,63 @@
+import json
+import shutil
 import xml.etree.ElementTree as ET
 from pathlib import Path
 from xml.dom import minidom
 
 import numpy as np
 import trimesh
+from tqdm import tqdm
 
 from happypose.toolbox.datasets.object_dataset import RigidObjectDataset
 
 
 def convert_rigid_body_dataset_to_urdfs(
-    rb_ds: RigidObjectDataset, urdf_dir: Path, texture_size=(1024, 1024)
+    rb_ds: RigidObjectDataset,
+    urdf_dir: Path,
+    texture_size=(1024, 1024),
+    override=True,
+    label2objname_file_name: str = "objname2label.json",
 ):
     """
-    Converts a RigidObjectDataset in a bullet renderer compatible directory with .obj and urdf files.
+    Converts a RigidObjectDataset into a directory of urdf files with structure:
+
+    urdf_dir/<label2objname_file_name>.json
+    urdf_dir/obj_000001/obj_000001.mtl
+                        obj_000001.obj
+                        obj_000001_texture.png
+                        obj_000001.urdf
+    urdf_dir/obj_000002/obj_000002.mtl
+                        obj_000002.obj
+                        obj_000002_texture.png
+                        obj_000002.urdf
+
+    <label2objname_file_name>.json: stores a map between object file names (e.g. obj_000002) and
+    object labels used in happypose (e.g. the detector may output "ycbv-obj_000002")
     """
+    if override and urdf_dir.exists():
+        shutil.rmtree(urdf_dir, ignore_errors=True)
     urdf_dir.mkdir(exist_ok=True, parents=True)
 
-    for obj in rb_ds.list_objects:
-        out_dir = urdf_dir / obj.mesh_path.with_suffix("").name
-        out_dir.mkdir(exist_ok=True)
-        obj_path = out_dir / obj.mesh_path.with_suffix(".obj").name
-        urdf_path = obj_path.with_suffix(".urdf")
+    objname2label = {}
+    for obj in tqdm(rb_ds.list_objects):
+        objname = obj.mesh_path.with_suffix("").name  # e.g. "obj_000002"
+        objname2label[objname] = obj.label  # e.g. obj_000002 -> ycbv-obj_000002
+        # Create object folder
+        obj_urdf_dir = urdf_dir / objname
+        obj_urdf_dir.mkdir(exist_ok=True)  # urdf_dir/obj_000002/ created
+        # Convert mesh from ply to obj
+        obj_path = (obj_urdf_dir / objname).with_suffix(".obj")
         if obj.mesh_path.suffix == ".ply":
-            if obj_path.exists():
-                obj_path.unlink()
             ply_to_obj(obj.mesh_path, obj_path, texture_size)
         else:
             ValueError(f"{obj.mesh_path.suffix} file type not supported")
+        # Create a .urdf file associated to the .obj file
+        urdf_path = obj_path.with_suffix(".urdf")
         obj_to_urdf(obj_path, urdf_path)
 
+    with open(urdf_dir / label2objname_file_name, "w") as fp:
+        json.dump(objname2label, fp)
+
 
 def ply_to_obj(ply_path: Path, obj_path: Path, texture_size=None):
     assert obj_path.suffix == ".obj"

happypose/pose_estimators/cosypose/cosypose/models/pose.py

@@ -24,6 +24,8 @@
     compute_rotation_matrix_from_quaternions,
 )
 from happypose.toolbox.renderer import Panda3dLightData
+from happypose.toolbox.renderer.bullet_batch_renderer import BulletBatchRenderer
+from happypose.toolbox.renderer.panda3d_batch_renderer import Panda3dBatchRenderer
 
 logger = get_logger(__name__)
 
@@ -128,21 +130,32 @@ def forward(self, images, K, labels, TCO, n_iterations=1):
                 labels,
             )
 
-            ambient_light = Panda3dLightData(
-                light_type="ambient",
-                color=(1.0, 1.0, 1.0, 1.0),
-            )
-            light_datas = [[ambient_light] for _ in range(len(labels))]
-
-            renders = self.renderer.render(
-                labels=labels,
-                TCO=TCO_input,
-                K=K_crop,
-                resolution=self.render_size,
-                light_datas=light_datas,
-            )
-            renders = renders.rgbs
-            x = torch.cat((images_crop, renders), dim=1)
+            if isinstance(self.renderer, Panda3dBatchRenderer):
+                ambient_light = Panda3dLightData(
+                    light_type="ambient",
+                    color=(1.0, 1.0, 1.0, 1.0),
+                )
+                light_datas = [[ambient_light] for _ in range(len(labels))]
+
+                renders = self.renderer.render(
+                    labels=labels,
+                    TCO=TCO_input,
+                    K=K_crop,
+                    resolution=self.render_size,
+                    light_datas=light_datas,
+                )
+            elif isinstance(self.renderer, BulletBatchRenderer):
+                renders = self.renderer.render(
+                    labels=labels,
+                    TCO=TCO_input,
+                    K=K_crop,
+                    resolution=self.render_size,
+                )
+            else:
+                raise ValueError(
+                    f"Renderer of type {type(self.renderer)} not supported"
+                )
+            x = torch.cat((images_crop, renders.rgbs), dim=1)
 
             model_outputs = self.net_forward(x)
 
@@ -158,7 +171,7 @@ def forward(self, images, K, labels, TCO, n_iterations=1):
             }
 
             outputs[f"iteration={n+1}"] = PosePredictorOutputCosypose(
-                renders=renders,
+                renders=renders.rgbs,
                 images_crop=images_crop,
                 TCO_input=TCO_input,
                 TCO_output=TCO_output,
@@ -176,7 +189,7 @@ def forward(self, images, K, labels, TCO, n_iterations=1):
                 self.tmp_debug.update(
                     images=images,
                     images_crop=images_crop,
-                    renders=renders,
+                    renders=renders.rgbs,
                 )
                 path = DEBUG_DATA_DIR / f"debug_iter={n+1}.pth.tar"
                 logger.info(f"Wrote debug data: {path}")

happypose/pose_estimators/cosypose/cosypose/scripts/convert_bop_ds_to_urdf.py

@@ -0,0 +1,43 @@
+import argparse
+from pathlib import Path
+
+from happypose.pose_estimators.cosypose.cosypose.libmesh.urdf_utils import (
+    convert_rigid_body_dataset_to_urdfs,
+)
+from happypose.toolbox.datasets.datasets_cfg import make_object_dataset
+
+
+def convert_bop_dataset_to_urdfs(
+    obj_ds_name: str, urdf_dir: Path, texture_size=(1024, 1024), override=True
+):
+    obj_dataset = make_object_dataset(obj_ds_name)
+    urdf_ds_dir = urdf_dir / obj_ds_name
+    convert_rigid_body_dataset_to_urdfs(
+        obj_dataset, urdf_ds_dir, texture_size, override
+    )
+
+
+def main():
+    parser = argparse.ArgumentParser("3D ply object ds_name -> pybullet URDF converter")
+    parser.add_argument(
+        "--ds_name",
+        default="",
+        type=str,
+        help="Bop dataset model name: ycbv, tless.cad, etc.",
+    )
+    parser.add_argument(
+        "--override",
+        default=True,
+        type=bool,
+        help="If true, erases previous content of urdf/<ds_name>.",
+    )
+    args = parser.parse_args()
+
+    from happypose.pose_estimators.cosypose.cosypose.config import LOCAL_DATA_DIR
+
+    urdf_dir = LOCAL_DATA_DIR / "urdfs"
+    convert_bop_dataset_to_urdfs(args.ds_name, urdf_dir, args.override)
+
+
+if __name__ == "__main__":
+    main()

happypose/pose_estimators/cosypose/cosypose/utils/cosypose_wrapper.py

@@ -29,7 +29,6 @@
 from happypose.toolbox.datasets.object_dataset import RigidObjectDataset
 from happypose.toolbox.inference.utils import load_detector
 from happypose.toolbox.lib3d.rigid_mesh_database import MeshDataBase
-from happypose.toolbox.renderer.panda3d_batch_renderer import Panda3dBatchRenderer
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
@@ -41,12 +40,24 @@ def __init__(
         object_dataset: Union[None, RigidObjectDataset] = None,
         n_workers=8,
         gpu_renderer=False,
+        renderer_type: str = "panda3d",
     ) -> None:
+        """
+        inputs:
+        - dataset_name: hope|tless|ycbv
+        - object_dataset: None or already existing rigid object dataset. If None, will use dataset_name
+        to build one
+        - n_workers: how many processes will be spun in the batch renderer
+        - renderer_type: 'panda3d'|'bullet'
+        """
+
         self.dataset_name = dataset_name
         self.object_dataset = object_dataset
-        self.detector, self.pose_predictor = self.get_model(dataset_name, n_workers)
+        self.detector, self.pose_predictor = self.get_model(
+            dataset_name, n_workers, renderer_type
+        )
 
-    def get_model(self, dataset_name, n_workers):
+    def get_model(self, dataset_name, n_workers, renderer_type):
         # load models
         if dataset_name == "hope":
             # HOPE setup
@@ -77,9 +88,7 @@ def get_model(self, dataset_name, n_workers):
             raise ValueError(msg)
         detector = load_detector(detector_run_id)
         coarse_model, refiner_model = self.load_pose_models(
-            coarse_run_id=coarse_run_id,
-            refiner_run_id=refiner_run_id,
-            n_workers=n_workers,
+            coarse_run_id, refiner_run_id, n_workers, renderer_type
         )
 
         pose_estimator = PoseEstimator(
@@ -89,14 +98,37 @@ def get_model(self, dataset_name, n_workers):
         )
         return detector, pose_estimator
 
-    def load_pose_models(self, coarse_run_id, refiner_run_id, n_workers):
+    def load_pose_models(
+        self,
+        coarse_run_id: str,
+        refiner_run_id: str,
+        n_workers: int,
+        renderer_type: str = "panda3d",
+    ):
         if self.object_dataset is None:
             self.object_dataset = make_object_dataset(self.dataset_name)
-        renderer = Panda3dBatchRenderer(
-            self.object_dataset,
-            n_workers=n_workers,
-            preload_cache=False,
-        )
+        if renderer_type == "panda3d":
+            from happypose.toolbox.renderer.panda3d_batch_renderer import (
+                Panda3dBatchRenderer,
+            )
+
+            renderer = Panda3dBatchRenderer(
+                self.object_dataset,
+                n_workers=n_workers,
+            )
+        elif renderer_type == "bullet":
+            from happypose.toolbox.renderer.bullet_batch_renderer import (
+                BulletBatchRenderer,
+            )
+
+            renderer = BulletBatchRenderer(
+                self.object_dataset,
+                n_workers=n_workers,
+                gpu_renderer=torch.cuda.is_available(),
+            )
+        else:
+            raise ValueError(f"Renderer type {renderer_type} not supported")
+
         mesh_db = MeshDataBase.from_object_ds(self.object_dataset)
         mesh_db_batched = mesh_db.batched().to(device)