[pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

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

@@ -203,10 +203,19 @@ def make_urdf_dataset(ds_name):
         return dataset
 
     # BOP
-    if ds_name in {"tless.cad", "tless.reconst", "ycbv", "hb", "icbin", "itodd", "lm", "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]

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

@@ -25,6 +25,7 @@ def __getitem__(self, idx):
     def __len__(self):
         return len(self.index)
 
+
 class OneUrdfDataset:
     def __init__(self, urdf_path, label, scale=1.0):
         index = [

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

@@ -8,40 +8,40 @@
 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)):
+def convert_rigid_body_dataset_to_urdfs(
+    rb_ds: RigidObjectDataset, urdf_dir: Path, texture_size=(1024, 1024)
+):
     """
     Converts a RigidObjectDataset in a bullet renderer compatible directory with .obj and urdf files.
     """
     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 = 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
+        obj_path = out_dir / obj.mesh_path.with_suffix(".obj").name
         urdf_path = obj_path.with_suffix(".urdf")
-        if obj.mesh_path.suffix == '.ply':
+        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')
+        else:
+            ValueError(f"{obj.mesh_path.suffix} file type not supported")
         obj_to_urdf(obj_path, urdf_path)
 
 
 def ply_to_obj(ply_path, obj_path, texture_size=None):
     mesh = trimesh.load(ply_path)
-    obj_label = obj_path.with_suffix('').name
+    obj_label = obj_path.with_suffix("").name
 
     # adapt materials according to previous example meshes
-    mesh.visual.material.ambient = np.array([51,51,51,255], dtype=np.uint8)
-    mesh.visual.material.diffuse = np.array([255,255,255,255], dtype=np.uint8)
-    mesh.visual.material.specular = np.array([255,255,255,255], dtype=np.uint8)
-    mesh.visual.material.name = obj_label + '_texture'
+    mesh.visual.material.ambient = np.array([51, 51, 51, 255], dtype=np.uint8)
+    mesh.visual.material.diffuse = np.array([255, 255, 255, 255], dtype=np.uint8)
+    mesh.visual.material.specular = np.array([255, 255, 255, 255], dtype=np.uint8)
+    mesh.visual.material.name = obj_label + "_texture"
 
     # print(mesh.visual.uv)
-    kwargs_export = {
-        'mtl_name': f'{obj_label}.mtl'
-    }
+    kwargs_export = {"mtl_name": f"{obj_label}.mtl"}
     _ = mesh.export(obj_path, **kwargs_export)
 
 

happypose/pose_estimators/cosypose/cosypose/scripts/convert_models_to_urdf.py

@@ -12,7 +12,9 @@
 )
 
 
-def convert_bop_dataset_to_urdfs(obj_ds_name: str, urdf_dir: Path, texture_size=(1024, 1024)):
+def convert_bop_dataset_to_urdfs(
+    obj_ds_name: str, urdf_dir: Path, texture_size=(1024, 1024)
+):
     obj_dataset = make_object_dataset(obj_ds_name)
     urdf_dir.mkdir(exist_ok=True, parents=True)
     for n in tqdm(range(len(obj_dataset))):
@@ -31,7 +33,8 @@ def main():
     args = parser.parse_args()
 
     from happypose.pose_estimators.cosypose.cosypose.config import LOCAL_DATA_DIR
-    urdf_dir = LOCAL_DATA_DIR / 'urdf'
+
+    urdf_dir = LOCAL_DATA_DIR / "urdf"
     convert_bop_dataset_to_urdfs(urdf_dir, args.models)
 
 

happypose/pose_estimators/cosypose/cosypose/simulator/body.py

@@ -17,6 +17,7 @@ def parse_pose_args(pose_args):
         raise ValueError
     return pose
 
+
 class Body:
     def __init__(self, body_id, scale=1.0, client_id=0):
         self._body_id = body_id

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

@@ -7,15 +7,16 @@
 import torch
 import transforms3d
 
+from happypose.pose_estimators.cosypose.cosypose.datasets.datasets_cfg import (
+    make_urdf_dataset,
+)
 from happypose.toolbox.lib3d.rotations import euler2quat
 from happypose.toolbox.lib3d.transform import Transform
 from happypose.toolbox.lib3d.transform_ops import invert_transform_matrices
 from happypose.toolbox.renderer.bullet_scene_renderer import (  # noqa: E501
     BulletSceneRenderer,
 )
-from happypose.pose_estimators.cosypose.cosypose.datasets.datasets_cfg import (
-    make_urdf_dataset,
-)
+
 from .plotter import Plotter
 
 
@@ -78,10 +79,7 @@ def make_scene_renderings(
 ):
     urdf_ds = make_urdf_dataset([urdf_ds_name, "camera"])
     renderer = BulletSceneRenderer(
-        urdf_ds,
-        background_color=background_color,
-        gui=gui,
-        gpu_renderer=False
+        urdf_ds, background_color=background_color, gui=gui, gpu_renderer=False
     )
     urdf_ds = renderer.body_cache.urdf_ds
 

happypose/pose_estimators/megapose/models/pose_rigid.py

@@ -430,7 +430,7 @@ def render_images_multiview(
             resolution=self.render_size,
             render_depth=self.render_depth,
             render_binary_mask=False,
-            render_normals=self.render_normals
+            render_normals=self.render_normals,
         )
 
         cat_list = []

happypose/toolbox/lib3d/transform.py

@@ -53,16 +53,18 @@ def __init__(
                 self._T = arg_T
             elif isinstance(arg_T, np.ndarray):
                 assert arg_T.shape == (4, 4)
-                assert arg_T[3,:3].sum() == 0.0 and arg_T[3,3] == 1.0
+                assert arg_T[3, :3].sum() == 0.0 and arg_T[3, 3] == 1.0
                 self._T = pin.SE3(arg_T)
             elif isinstance(arg_T, torch.Tensor):
                 T = arg_T.detach().cpu().numpy().copy()
                 assert T.shape == (4, 4)
-                assert T[3,:3].sum() == 0.0 and T[3,3] == 1.0
+                assert T[3, :3].sum() == 0.0 and T[3, 3] == 1.0
                 self._T = pin.SE3(T)
             else:
-                raise ValueError(f"""Transform contructor: if 1 argument, should be a Transform, 
-                                 pin.SE3, numpy array, or torch Tensor, not {type(arg_T)}""")
+                raise ValueError(
+                    f"""Transform contructor: if 1 argument, should be a Transform,
+                                 pin.SE3, numpy array, or torch Tensor, not {type(arg_T)}"""
+                )
 
         elif len(args) == 2:
             rotation, translation = args
@@ -76,9 +78,11 @@ def __init__(
                 else:
                     rotation_np = rotation
             else:
-                raise ValueError('Transform contructor: if 2 argument, rotation should be either a pin.Quaternion, \
+                raise ValueError(
+                    "Transform contructor: if 2 argument, rotation should be either a pin.Quaternion, \
                                  a 4-d tuple|list representing xyzw quaternion, \
-                                 or a np.array|torch.Tensor (either 4d xyzw quaternion or 3x3 SO(3) matrix)')
+                                 or a np.array|torch.Tensor (either 4d xyzw quaternion or 3x3 SO(3) matrix)"
+                )
 
             if rotation_np.size == 4:
                 quaternion_xyzw = rotation_np.flatten().tolist()
@@ -131,4 +135,4 @@ def matrix(self) -> np.ndarray:
 
     @staticmethod
     def Identity():
-        return Transform(pin.SE3.Identity())
+        return Transform(pin.SE3.Identity())

happypose/toolbox/renderer/bullet_batch_renderer.py

@@ -1,5 +1,5 @@
 import multiprocessing
-from typing import Dict, List, Union
+from typing import List, Union
 
 import numpy as np
 import torch
@@ -10,7 +10,11 @@
 from happypose.toolbox.renderer.types import BatchRenderOutput, WorkerRenderOutput
 
 
-def init_renderer(asset_dataset: Union[RigidObjectDataset,UrdfDataset], preload=True, gpu_renderer=True):
+def init_renderer(
+    asset_dataset: Union[RigidObjectDataset, UrdfDataset],
+    preload=True,
+    gpu_renderer=True,
+):
     renderer = BulletSceneRenderer(
         asset_dataset=asset_dataset,
         preload_cache=preload,
@@ -33,7 +37,7 @@ def worker_loop(
         render_args = in_queue.get()
         if render_args is None:
             return
-        
+
         obj_infos = render_args["obj_infos"]
         cam_infos = render_args["cam_infos"]
         render_depth = render_args["render_depth"]
@@ -51,8 +55,14 @@ def worker_loop(
                 render_binary_mask=render_binary_mask,
             )
             rgbs = np.stack([ren.rgb for ren in renderings])
-            depth = np.stack([ren.depth for ren in renderings]) if render_depth else None
-            binary_mask = np.stack([ren.binary_mask for ren in renderings]) if render_binary_mask else None
+            depth = (
+                np.stack([ren.depth for ren in renderings]) if render_depth else None
+            )
+            binary_mask = (
+                np.stack([ren.binary_mask for ren in renderings])
+                if render_binary_mask
+                else None
+            )
         else:
             w, h = cam_infos[0]["resolution"]
             rgbs = np.zeros((1, h, w, 3), dtype=np.uint8)
@@ -70,24 +80,27 @@ def worker_loop(
 
 
 class BulletBatchRenderer:
-    def __init__(self, 
-                 asset_dataset: Union[RigidObjectDataset,UrdfDataset], 
-                 n_workers=8, 
-                 preload_cache=True, 
-                 gpu_renderer=True):
+    def __init__(
+        self,
+        asset_dataset: Union[RigidObjectDataset, UrdfDataset],
+        n_workers=8,
+        preload_cache=True,
+        gpu_renderer=True,
+    ):
         self.asset_dataset = asset_dataset
         self.n_workers = n_workers
         self.init_plotters(preload_cache, gpu_renderer)
         self.gpu_renderer = gpu_renderer
 
-    def render(self, 
-               labels: List[str], 
-               TCO: torch.Tensor, 
-               K: torch.Tensor, 
-               resolution=(240, 320), 
-               render_depth: bool=False,
-               render_binary_mask: bool=False
-               ):
+    def render(
+        self,
+        labels: List[str],
+        TCO: torch.Tensor,
+        K: torch.Tensor,
+        resolution=(240, 320),
+        render_depth: bool = False,
+        render_binary_mask: bool = False,
+    ):
         TCO = torch.as_tensor(TCO).detach()
         TOC = invert_transform_matrices(TCO).cpu().numpy()
         K = torch.as_tensor(K).cpu().numpy()
@@ -121,7 +134,7 @@ def render(self,
                 self.in_queue.put(render_args)
             else:
                 renderings = self.plotters[0].render_scene(**render_args)
-                # by definition, each "scene" in batch rendering corresponds to 1 camera, 1 object 
+                # by definition, each "scene" in batch rendering corresponds to 1 camera, 1 object
                 # -> retrieves the first and only rendering
                 renderings_ = renderings[0]
 
@@ -170,11 +183,13 @@ def render(self,
             else:
                 depths = torch.stack(list_depths)
             depths = depths.float().permute(0, 3, 1, 2)
-        
+
         if render_binary_mask:
             assert list_binary_masks[0] is not None
             if torch.cuda.is_available():
-                binary_masks = torch.stack(list_binary_masks).pin_memory().cuda(non_blocking=True)
+                binary_masks = (
+                    torch.stack(list_binary_masks).pin_memory().cuda(non_blocking=True)
+                )
             else:
                 binary_masks = torch.stack(list_binary_masks)
             binary_masks = binary_masks.permute(0, 3, 1, 2)

happypose/toolbox/renderer/bullet_scene_renderer.py

@@ -25,7 +25,7 @@
 class BulletSceneRenderer(BaseScene):
     def __init__(
         self,
-        asset_dataset: Union[RigidObjectDataset,UrdfDataset],
+        asset_dataset: Union[RigidObjectDataset, UrdfDataset],
         preload_cache=False,
         background_color=(0, 0, 0),
         gpu_renderer=True,
@@ -35,14 +35,16 @@ def __init__(
             self.urdf_ds = asset_dataset
         elif isinstance(asset_dataset, RigidObjectDataset):
             # Build urdfs files from RigidObjectDataset
-            ds_name = 'tmp'
+            ds_name = "tmp"
             urdf_dir = LOCAL_DATA_DIR / "urdfs" / ds_name
             convert_rigid_body_dataset_to_urdfs(asset_dataset, urdf_dir)
             self.urdf_ds = UrdfDataset(urdf_dir)
             # TODO: BodyCache assumes unique scale for all objects (true for bop datasets)
             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' )
+            raise TypeError(
+                f"asset_dataset of type {type(asset_dataset)} should be either UrdfDataset or RigidObjectDataset"
+            )
         self.connect(gpu_renderer=gpu_renderer, gui=gui)
         self.body_cache = BodyCache(self.urdf_ds, self.client_id)
         if preload_cache:
@@ -64,14 +66,14 @@ def setup_scene(self, obj_infos):
                 )
         return bodies
 
-    def render_images(self, 
-                      cam_infos: Dict, 
-                      render_depth: bool=False,
-                      render_binary_mask: bool=False
-                      ) -> List[CameraRenderingData]:
+    def render_images(
+        self,
+        cam_infos: Dict,
+        render_depth: bool = False,
+        render_binary_mask: bool = False,
+    ) -> List[CameraRenderingData]:
         cam_renderings = []
         for cam_info in cam_infos:
-
             K = cam_info["K"]
             TWC = Transform(cam_info["TWC"])
             resolution = cam_info["resolution"]
@@ -103,14 +105,17 @@ def render_images(self,
                 rendering.depth = z_e[..., np.newaxis]  # (h,w,1)
 
             cam_renderings.append(rendering)
-            
+
         return cam_renderings
 
-    def render_scene(self, 
-                     obj_infos: Dict, 
-                     cam_infos: Dict, 
-                     render_depth=False,
-                     render_binary_mask=False
-                     ) -> List[CameraRenderingData]:
+    def render_scene(
+        self,
+        obj_infos: Dict,
+        cam_infos: Dict,
+        render_depth=False,
+        render_binary_mask=False,
+    ) -> List[CameraRenderingData]:
         self.setup_scene(obj_infos)
-        return self.render_images(cam_infos, render_depth=render_depth, render_binary_mask=render_binary_mask)
+        return self.render_images(
+            cam_infos, render_depth=render_depth, render_binary_mask=render_binary_mask
+        )