Merge pull request #130 from agimus-project/122-renderers-refactoring

122 renderers refactoring

README.md

@@ -48,7 +48,7 @@ pip install .[pypi,cpu] --extra-index-url https://download.pytorch.org/whl/cpu
 - `gpu`: required to get pytorch GPU from PyPI (don't use this for CPU or with conda)
 - `evaluation`: installs bop_toolkit
 - `multiview`: installs cosypose c++ extension
-- `pypi`: install pinocchio & opncv from PyPI (don't use this with conda)
+- `pypi`: install pinocchio & opencv from PyPI (don't use this with conda)
 
 ## Create data directory
 

happypose/pose_estimators/cosypose/cosypose/datasets/datasets_cfg.py

@@ -11,7 +11,7 @@
 from .bop import BOPDataset, remap_bop_targets
 from .bop_object_datasets import BOPObjectDataset
 from .texture_dataset import TextureDataset
-from .urdf_dataset import BOPUrdfDataset, OneUrdfDataset
+from .urdf_dataset import OneUrdfDataset, UrdfDataset
 
 logger = get_logger(__name__)
 
@@ -158,7 +158,7 @@ def make_scene_dataset(ds_name, n_frames=None):
     return ds
 
 
-def make_object_dataset(ds_name):
+def make_object_dataset(ds_name: str):
     ds = None
     if ds_name == "tless.cad":
         ds = BOPObjectDataset(BOP_DS_DIR / "tless/models_cad")
@@ -203,27 +203,24 @@ def make_urdf_dataset(ds_name):
         return dataset
 
     # BOP
-    if ds_name == "tless.cad":
-        ds = BOPUrdfDataset(LOCAL_DATA_DIR / "urdfs" / "tless.cad")
-    elif ds_name == "tless.reconst":
-        ds = BOPUrdfDataset(LOCAL_DATA_DIR / "urdfs" / "tless.reconst")
-    elif ds_name == "ycbv":
-        ds = BOPUrdfDataset(LOCAL_DATA_DIR / "urdfs" / "ycbv")
-    elif ds_name == "hb":
-        ds = BOPUrdfDataset(LOCAL_DATA_DIR / "urdfs" / "hb")
-    elif ds_name == "icbin":
-        ds = BOPUrdfDataset(LOCAL_DATA_DIR / "urdfs" / "icbin")
-    elif ds_name == "itodd":
-        ds = BOPUrdfDataset(LOCAL_DATA_DIR / "urdfs" / "itodd")
-    elif ds_name == "lm":
-        ds = BOPUrdfDataset(LOCAL_DATA_DIR / "urdfs" / "lm")
-    elif ds_name == "tudl":
-        ds = BOPUrdfDataset(LOCAL_DATA_DIR / "urdfs" / "tudl")
+    if ds_name in {
+        "tless.cad",
+        "tless.reconst",
+        "ycbv",
+        "hb",
+        "icbin",
+        "itodd",
+        "lm",
+        "tudl",
+    }:
+        ds = UrdfDataset(LOCAL_DATA_DIR / "urdfs" / ds_name)
+        ds.index["scale"] = 0.001
 
     # Custom scenario
     elif "custom" in ds_name:
         scenario = ds_name.split(".")[1]
-        ds = BOPUrdfDataset(LOCAL_DATA_DIR / "scenarios" / scenario / "urdfs")
+        ds = UrdfDataset(LOCAL_DATA_DIR / "scenarios" / scenario / "urdfs")
+        ds.index["scale"] = 0.001
 
     elif ds_name == "camera":
         ds = OneUrdfDataset(ASSET_DIR / "camera/model.urdf", "camera")

happypose/pose_estimators/cosypose/cosypose/datasets/pose_dataset.py

@@ -5,7 +5,7 @@
 import torch
 
 from happypose.pose_estimators.cosypose.cosypose.config import LOCAL_DATA_DIR
-from happypose.pose_estimators.cosypose.cosypose.lib3d.transform_ops import invert_T
+from happypose.toolbox.lib3d.transform_ops import invert_transform_matrices
 
 from .augmentations import (
     CropResizeToAspectAugmentation,
@@ -125,7 +125,7 @@ def get_data(self, idx):
         obj = random.sample(objects_visible, k=1)[0]
         TWO = torch.as_tensor(obj["TWO"])
         TWC = torch.as_tensor(state["camera"]["TWC"])
-        TCO = invert_T(TWC) @ TWO
+        TCO = invert_transform_matrices(TWC) @ TWO
 
         data = PoseData(
             images=np.asarray(rgb),

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

@@ -26,12 +26,6 @@ def __len__(self):
         return len(self.index)
 
 
-class BOPUrdfDataset(UrdfDataset):
-    def __init__(self, ds_dir):
-        super().__init__(ds_dir)
-        self.index["scale"] = 0.001
-
-
 class OneUrdfDataset:
     def __init__(self, urdf_path, label, scale=1.0):
         index = [
@@ -44,17 +38,3 @@ def __len__(self):
 
     def __getitem__(self, idx):
         return self.index.iloc[idx]
-
-
-class UrdfMultiScaleDataset(UrdfDataset):
-    def __init__(self, urdf_path, label, scales=[]):
-        index = []
-        for scale in scales:
-            index.append(
-                {
-                    "urdf_path": urdf_path,
-                    "label": label + f"scale={scale:.3f}",
-                    "scale": scale,
-                },
-            )
-        self.index = pd.DataFrame(index)

happypose/pose_estimators/cosypose/cosypose/evaluation/data_utils.py

@@ -4,7 +4,7 @@
 import torch
 
 import happypose.pose_estimators.cosypose.cosypose.utils.tensor_collection as tc
-from happypose.pose_estimators.cosypose.cosypose.lib3d.transform_ops import invert_T
+from happypose.toolbox.lib3d.transform_ops import invert_transform_matrices
 
 
 def parse_obs_data(obs):
@@ -28,7 +28,7 @@ def parse_obs_data(obs):
             data[k] = torch.stack([torch.as_tensor(x).float() for x in v])
 
     data["infos"] = pd.DataFrame(data["infos"])
-    TCO = invert_T(TWC).unsqueeze(0) @ data["TWO"]
+    TCO = invert_transform_matrices(TWC).unsqueeze(0) @ data["TWO"]
 
     data = tc.PandasTensorCollection(
         infos=data["infos"],

happypose/pose_estimators/cosypose/cosypose/integrated/multiview_predictor.py

@@ -1,8 +1,8 @@
 import numpy as np
 import pandas as pd
+import torch
 
 import happypose.pose_estimators.cosypose.cosypose.utils.tensor_collection as tc
-from happypose.pose_estimators.cosypose.cosypose.lib3d.transform_ops import invert_T
 from happypose.pose_estimators.cosypose.cosypose.multiview.bundle_adjustment import (
     MultiviewRefinement,
     make_view_groups,
@@ -11,16 +11,21 @@
     multiview_candidate_matching,
 )
 from happypose.pose_estimators.cosypose.cosypose.utils.logging import get_logger
+from happypose.toolbox.lib3d.transform_ops import invert_transform_matrices
 
 logger = get_logger(__name__)
 
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
 
 class MultiviewScenePredictor:
     def __init__(self, mesh_db, n_sym=64, ba_aabb=True, ba_n_points=None):
-        self.mesh_db_ransac = mesh_db.batched(n_sym=n_sym, aabb=True).float()
-        self.mesh_db_ba = mesh_db.batched(
-            aabb=ba_aabb, resample_n_points=ba_n_points, n_sym=n_sym
-        ).float()
+        self.mesh_db_ransac = mesh_db.batched(n_sym=n_sym, aabb=True).to(device).float()
+        self.mesh_db_ba = (
+            mesh_db.batched(aabb=ba_aabb, resample_n_points=ba_n_points, n_sym=n_sym)
+            .to(device)
+            .float()
+        )
 
     def reproject_scene(self, objects, cameras):
         TCO_data = []
@@ -40,7 +45,7 @@ def reproject_scene(self, objects, cameras):
                 }
                 data_ = tc.PandasTensorCollection(
                     infos=pd.DataFrame(infos),
-                    poses=invert_T(cam.TWC) @ obj.TWO,
+                    poses=invert_transform_matrices(cam.TWC) @ obj.TWO,
                 )
                 TCO_data.append(data_)
         return tc.concatenate(TCO_data)
@@ -88,7 +93,7 @@ def predict_scene_state(
         predictions["cand_matched"] = candidates
 
         group_infos = make_view_groups(pairs_TC1C2)
-        candidates = candidates.merge_df(group_infos, on="view_id").cuda()
+        candidates = candidates.merge_df(group_infos, on="view_id").to(device)
 
         pred_objects, pred_cameras, pred_reproj = [], [], []
         pred_reproj_init = []

happypose/pose_estimators/cosypose/cosypose/lib3d/cosypose_ops.py

@@ -1,10 +1,10 @@
 import torch
 
-from .rotations import (
+from happypose.toolbox.lib3d.rotations import (
     compute_rotation_matrix_from_ortho6d,
     compute_rotation_matrix_from_quaternions,
 )
-from .transform_ops import transform_pts
+from happypose.toolbox.lib3d.transform_ops import transform_pts
 
 
 def l1(diff):

happypose/pose_estimators/cosypose/cosypose/lib3d/symmetries.py

@@ -1,7 +1,7 @@
 import numpy as np
 
-from .rotations import euler2quat
-from .transform import Transform
+from happypose.toolbox.lib3d.rotations import euler2quat
+from happypose.toolbox.lib3d.transform import Transform
 
 
 def make_bop_symmetries(dict_symmetries, n_symmetries_continuous=8, scale=0.001):