Adapting slightly BulletBatchRenderer to new API, needs refactoring

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

@@ -101,13 +101,27 @@ def forward(self, images, K, labels, TCO, n_iterations=1):
             TCO_input = TCO_input.detach()
             images_crop, K_crop, boxes_rend, boxes_crop = self.crop_inputs(images, K, TCO_input, 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,)
+            # HACK: ugly solution, API of the renderer should be uniformized 
+            if type(self.renderer).__name__ == '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 type(self.renderer).__name__ == 'BulletBatchRenderer':
+                # labels: ['ycbv-obj_000006', 'ycbv-obj_000009', 'ycbv-obj_000003', 'ycbv-obj_000013']
+                # but BulletBatchRenderer expects ['obj_{0000:id}']
+                labels_render = [l.split('-')[1] for l in labels]
+                renders = self.renderer.render(obj_infos=[dict(name=l) for l in labels_render],
+                                            TCO=TCO_input,
+                                            K=K_crop, resolution=self.render_size)
+
+            else:
+                ValueError(f'{type(self.renderer).__name__} renderer type not supported by cosypose')
+
             renders = renders.rgbs
             x = torch.cat((images_crop, renders), dim=1)
 

happypose/pose_estimators/cosypose/cosypose/rendering/bullet_batch_renderer.py

@@ -1,6 +1,9 @@
 import torch
 import numpy as np
 import multiprocessing
+from dataclasses import dataclass
+from typing import Optional
+
 
 from happypose.pose_estimators.cosypose.cosypose.lib3d.transform_ops import invert_T
 from .bullet_scene_renderer import BulletSceneRenderer
@@ -38,6 +41,21 @@ def worker_loop(worker_id, in_queue, out_queue, object_set, preload=True, gpu_re
         out_queue.put((kwargs['data_id'], images, depth))
 
 
+# same API as Panda3d
+# TODO: refactor to happypose_toolbox 
+@dataclass
+class BatchRenderOutput:
+    """
+    rgb: (bsz, 3, h, w) float, values in [0, 1]
+    normals: (bsz, 3, h, w) float, values in [0, 1]
+    depth: (bsz, 1, h, w) float, in meters.
+    """
+
+    rgbs: torch.Tensor
+    normals: Optional[torch.Tensor] = None
+    depths: Optional[torch.Tensor] = None
+
+
 class BulletBatchRenderer:
     def __init__(self, object_set, n_workers=8, preload_cache=True, gpu_renderer=True):
         self.object_set = object_set
@@ -93,9 +111,9 @@ def render(self, obj_infos, TCO, K, resolution=(240, 320), render_depth=False):
             else:
                 depths = torch.as_tensor(np.stack(depths, axis=0))
             depths = depths.float()
-            return images, depths
+            return BatchRenderOutput(images, depths)
         else:
-            return images
+            return BatchRenderOutput(images)
 
     def init_plotters(self, preload_cache, gpu_renderer):
         self.plotters = []

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

@@ -13,36 +13,55 @@
 - refactor hardcoded model weight checkpoints
 """
 
-from PIL import Image
-import numpy as np
+import argparse
 from copy import deepcopy
 from pathlib import Path
-import yaml
-import torch
-import argparse
-import pandas as pd
 
-# from happypose.pose_estimators.cosypose.cosypose.datasets.datasets_cfg import make_scene_dataset, make_object_dataset
-from happypose.toolbox.datasets.datasets_cfg import make_scene_dataset, make_object_dataset
+import numpy as np
+import pandas as pd
+import torch
+import yaml
+from PIL import Image
 
-# Pose estimator
-from happypose.toolbox.lib3d.rigid_mesh_database import MeshDataBase
-from happypose.pose_estimators.cosypose.cosypose.training.pose_models_cfg import create_model_refiner, create_model_coarse
-from happypose.pose_estimators.cosypose.cosypose.training.pose_models_cfg import check_update_config as check_update_config_pose
-# Detection
-from happypose.pose_estimators.cosypose.cosypose.training.detector_models_cfg import create_model_detector
-from happypose.pose_estimators.cosypose.cosypose.training.detector_models_cfg import check_update_config as check_update_config_detector
+from happypose.pose_estimators.cosypose.cosypose.config import EXP_DIR, RESULTS_DIR
+from happypose.pose_estimators.cosypose.cosypose.evaluation.pred_runner.bop_predictions import (
+    BopPredictionRunner,
+)
 from happypose.pose_estimators.cosypose.cosypose.integrated.detector import Detector
-
-from happypose.pose_estimators.cosypose.cosypose.evaluation.pred_runner.bop_predictions import BopPredictionRunner
 from happypose.pose_estimators.cosypose.cosypose.integrated.pose_estimator import (
     PoseEstimator,
 )
-from happypose.pose_estimators.cosypose.cosypose.utils.distributed import get_tmp_dir, get_rank
-from happypose.pose_estimators.cosypose.cosypose.utils.distributed import init_distributed_mode
 
-from happypose.pose_estimators.cosypose.cosypose.config import EXP_DIR, RESULTS_DIR
+# Detection
+from happypose.pose_estimators.cosypose.cosypose.training.detector_models_cfg import (
+    check_update_config as check_update_config_detector,
+)
+from happypose.pose_estimators.cosypose.cosypose.training.detector_models_cfg import (
+    create_model_detector,
+)
+from happypose.pose_estimators.cosypose.cosypose.training.pose_models_cfg import (
+    check_update_config as check_update_config_pose,
+)
+from happypose.pose_estimators.cosypose.cosypose.training.pose_models_cfg import (
+    create_model_coarse,
+    create_model_refiner,
+)
+from happypose.pose_estimators.cosypose.cosypose.utils.distributed import (
+    get_rank,
+    get_tmp_dir,
+    init_distributed_mode,
+)
 
+# from happypose.pose_estimators.cosypose.cosypose.datasets.datasets_cfg import make_scene_dataset, make_object_dataset
+from happypose.toolbox.datasets.datasets_cfg import (
+    make_object_dataset,
+    make_scene_dataset,
+)
+
+# Pose estimator
+from happypose.toolbox.lib3d.rigid_mesh_database import MeshDataBase
+
+from happypose.pose_estimators.cosypose.cosypose.rendering.bullet_batch_renderer import BulletBatchRenderer
 from happypose.toolbox.renderer.panda3d_batch_renderer import Panda3dBatchRenderer
 
 
@@ -75,10 +94,11 @@ def make_object_dataset(example_dir: Path) -> RigidObjectDataset:
 
 
 class CosyPoseWrapper:
-    def __init__(self, dataset_name, n_workers=8, gpu_renderer=False) -> None:
+    def __init__(self, dataset_name, n_workers=8, renderer_name='bullet') -> None:
+        self.renderer_name = renderer_name
         self.dataset_name = dataset_name
         super().__init__()
-        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_name)
 
     @staticmethod
     def load_detector(run_id, ds_name):
@@ -98,21 +118,23 @@ def load_detector(run_id, ds_name):
         return model
 
     @staticmethod
-    def load_pose_models(coarse_run_id, refiner_run_id, n_workers):
+    def load_pose_models(coarse_run_id, refiner_run_id, n_workers, renderer_name):
         run_dir = EXP_DIR / coarse_run_id
 
         # cfg = yaml.load((run_dir / 'config.yaml').read_text(), Loader=yaml.FullLoader)
         cfg = yaml.load((run_dir / 'config.yaml').read_text(), Loader=yaml.UnsafeLoader)
         cfg = check_update_config_pose(cfg)
-        # object_ds = BOPObjectDataset(BOP_DS_DIR / 'tless/models_cad')
-        #object_ds = make_object_dataset(cfg.object_ds_name)
-        #mesh_db = MeshDataBase.from_object_ds(object_ds)
-        #renderer = BulletBatchRenderer(object_set=cfg.urdf_ds_name, n_workers=n_workers, gpu_renderer=gpu_renderer)
-        #
-
+        # object_dataset = BOPObjectDataset(BOP_DS_DIR / 'tless/models_cad')
+        # object_dataset = make_object_dataset(cfg.object_ds_name)
         object_dataset = make_object_dataset("ycbv")
         mesh_db = MeshDataBase.from_object_ds(object_dataset)
-        renderer = Panda3dBatchRenderer(object_dataset, n_workers=n_workers, preload_cache=False)
+
+        if renderer_name == 'bullet':
+            renderer = BulletBatchRenderer(object_set=cfg.urdf_ds_name, n_workers=n_workers, gpu_renderer=device.type == 'cuda')
+        elif renderer_name == 'panda':
+            renderer = Panda3dBatchRenderer(object_dataset, n_workers=n_workers, preload_cache=False)
+        else:
+            raise ValueError(f'{renderer_name} is not supported, choose between "bullet" and "panda"')
         mesh_db_batched = mesh_db.batched().to(device)
 
         def load_model(run_id):
@@ -138,7 +160,7 @@ def load_model(run_id):
         return coarse_model, refiner_model, mesh_db
 
     @staticmethod
-    def get_model(dataset_name, n_workers):
+    def get_model(dataset_name, n_workers, renderer_name):
         # load models
         if dataset_name == 'tless':
             # TLESS setup
@@ -160,7 +182,7 @@ def get_model(dataset_name, n_workers):
             raise ValueError(f"Not prepared for {dataset_name} dataset")
         detector = CosyPoseWrapper.load_detector(detector_run_id, dataset_name)
         coarse_model, refiner_model , mesh_db = CosyPoseWrapper.load_pose_models(
-            coarse_run_id=coarse_run_id, refiner_run_id=refiner_run_id, n_workers=n_workers
+            coarse_run_id=coarse_run_id, refiner_run_id=refiner_run_id, n_workers=n_workers, renderer_name
         )
 
         pose_estimator = PoseEstimator(