Merge branch 'cosy-eval-adapt' of github.com:agimus-project/happypose into cosy-eval-adapt

Author: kzorina

docs/book/cosypose/evaluate.md

@@ -14,4 +14,4 @@ python -m happypose.pose_estimators.cosypose.cosypose.scripts.run_full_cosypose_
 ```
 
 To evaluate on other datasets, change ["ycbv.bop19"] to e.g. ["tless.bop19"].
-To evaluate on a collection of datasets, change "ycbv.bop19" to e.g. ["ycbv.bop19", "lmo.bop19", "tless.bop19"].
+To evaluate on a collection of datasets, change "ycbv.bop19" to e.g. ["ycbv.bop19", "lmo.bop19", "tless.bop19"].

happypose/pose_estimators/cosypose/cosypose/config.py

@@ -21,7 +21,9 @@
 LOCAL_DATA_DIR = Path(
     os.environ.get("HAPPYPOSE_DATA_DIR", Path(PROJECT_DIR) / "local_data"),
 )
-assert LOCAL_DATA_DIR.exists(), "Did you forget to set env variable 'HAPPYPOSE_DATA_DIR'?"
+assert (
+    LOCAL_DATA_DIR.exists()
+), "Did you forget to set env variable 'HAPPYPOSE_DATA_DIR'?"
 TEST_DATA_DIR = LOCAL_DATA_DIR
 DASK_LOGS_DIR = LOCAL_DATA_DIR / "dasklogs"
 SYNT_DS_DIR = LOCAL_DATA_DIR / "synt_datasets"

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

@@ -33,7 +33,11 @@
 from happypose.pose_estimators.megapose.evaluation.runner_utils import format_results
 
 # Pose estimator
-from happypose.toolbox.datasets.datasets_cfg import make_object_dataset, make_scene_dataset, get_obj_ds_info
+from happypose.toolbox.datasets.datasets_cfg import (
+    get_obj_ds_info,
+    make_object_dataset,
+    make_scene_dataset,
+)
 from happypose.toolbox.lib3d.rigid_mesh_database import MeshDataBase
 from happypose.toolbox.utils.distributed import get_rank, get_tmp_dir
 from happypose.toolbox.utils.logging import get_logger
@@ -44,13 +48,14 @@
 
 logger = get_logger(__name__)
 
+from happypose.pose_estimators.cosypose.cosypose.integrated.detector import 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.training.detector_models_cfg import (
     create_model_detector,
 )
-from happypose.pose_estimators.cosypose.cosypose.integrated.detector import Detector
+
 
 def load_detector(run_id, ds_name):
     run_dir = EXP_DIR / run_id
@@ -68,7 +73,9 @@ def load_detector(run_id, ds_name):
     return model
 
 
-def load_pose_models_cosypose(object_dataset, coarse_run_id, refiner_run_id, n_workers, renderer_type="panda3d"):
+def load_pose_models_cosypose(
+    object_dataset, coarse_run_id, refiner_run_id, n_workers, renderer_type="panda3d"
+):
     run_dir = EXP_DIR / coarse_run_id
     cfg = yaml.load((run_dir / "config.yaml").read_text(), Loader=yaml.UnsafeLoader)
     cfg = check_update_config_pose(cfg)
@@ -161,7 +168,7 @@ def run_eval(
 
     logger.info(f"Running eval on ds_name={cfg.ds_name} with setting={save_key}")
     # e.g. "ycbv.bop19" -> "ycbv"
-    ds_name_short = cfg.ds_name.split('.')[0]
+    ds_name_short = cfg.ds_name.split(".")[0]
 
     # Load the dataset
     ds_kwargs = {"load_depth": False}
@@ -195,15 +202,14 @@ def run_eval(
     assert cfg.coarse_run_id is not None
     assert cfg.refiner_run_id is not None
 
-
     object_ds = make_object_dataset(ds_name_short)
 
     coarse_model, refiner_model, mesh_db = load_pose_models_cosypose(
         object_ds,
         coarse_run_id=cfg.coarse_run_id,
         refiner_run_id=cfg.refiner_run_id,
         n_workers=cfg.inference.n_workers,
-        renderer_type=cfg.inference.renderer
+        renderer_type=cfg.inference.renderer,
     )
 
     renderer = refiner_model.renderer

happypose/pose_estimators/cosypose/cosypose/evaluation/prediction_runner.py

@@ -122,7 +122,7 @@ def run_inference_pipeline(
             data_TCO_init=None,
             n_coarse_iterations=1,
             n_refiner_iterations=4,
-            detection_th=0.0
+            detection_th=0.0,
         )
         time.time() - t
 

happypose/pose_estimators/cosypose/cosypose/scripts/run_full_cosypose_eval_new.py

@@ -10,7 +10,10 @@
 from happypose.pose_estimators.cosypose.cosypose.evaluation.evaluation import run_eval
 
 # MegaPose
-from happypose.pose_estimators.megapose.bop_config import PBR_DETECTORS, SYNT_REAL_DETECTORS
+from happypose.pose_estimators.megapose.bop_config import (
+    PBR_DETECTORS,
+    SYNT_REAL_DETECTORS,
+)
 from happypose.pose_estimators.megapose.config import (
     DEBUG_RESULTS_DIR,
     RESULTS_DIR,

happypose/pose_estimators/megapose/config.py

@@ -31,7 +31,9 @@
 LOCAL_DATA_DIR = Path(
     os.environ.get("HAPPYPOSE_DATA_DIR", Path(PROJECT_DIR) / "local_data"),
 )
-assert LOCAL_DATA_DIR.exists(), "Did you forget to set env variable 'HAPPYPOSE_DATA_DIR'?"
+assert (
+    LOCAL_DATA_DIR.exists()
+), "Did you forget to set env variable 'HAPPYPOSE_DATA_DIR'?"
 BOP_DS_DIR = LOCAL_DATA_DIR / "bop_datasets"
 NB_DATA_DIR = LOCAL_DATA_DIR / "notebook_data"
 SHAPENET_DIR = LOCAL_DATA_DIR / "shapenetcorev2"

happypose/toolbox/datasets/bop_scene_dataset.py

@@ -210,7 +210,7 @@ def __init__(
         use_raw_object_id: bool = False,
         allow_cache: bool = False,
         per_view_annotations: bool = False,
-        models_dir: str = "models"
+        models_dir: str = "models",
     ):
         assert ds_dir.exists(), "Dataset does not exists."
         self.ds_dir = ds_dir
@@ -247,7 +247,9 @@ def __init__(
             load_segmentation=True,
         )
         models_path = ds_dir / models_dir
-        assert models_path.exists(), f"models_dir={models_dir} might not be the correct object model dir name"
+        assert (
+            models_path.exists()
+        ), f"models_dir={models_dir} might not be the correct object model dir name"
         models_infos = json.loads((models_path / "models_info.json").read_text())
         self.all_labels = [f"obj_{int(obj_id):06d}" for obj_id in models_infos.keys()]
 

happypose/toolbox/datasets/datasets_cfg.py

@@ -91,7 +91,12 @@ def make_scene_dataset(
         ds = keep_bop19(ds)
     elif ds_name == "tless.bop19":
         ds_dir = BOP_DS_DIR / "tless"
-        ds = BOPDataset(ds_dir, split="test_primesense", label_format="tless-{label}", models_dir="models_cad")
+        ds = BOPDataset(
+            ds_dir,
+            split="test_primesense",
+            label_format="tless-{label}",
+            models_dir="models_cad",
+        )
         ds = keep_bop19(ds)
     elif ds_name == "tudl.bop19":
         ds_dir = BOP_DS_DIR / "tudl"

happypose/toolbox/renderer/bullet_scene_renderer.py

@@ -1,18 +1,15 @@
 from typing import Dict, List, Union
 
-import json
 import numpy as np
 import pybullet as pb
 
 from happypose.pose_estimators.cosypose.cosypose.config import LOCAL_DATA_DIR
-
 from happypose.pose_estimators.cosypose.cosypose.datasets.urdf_dataset import (
-    UrdfDataset,
     OneUrdfDataset,
+    UrdfDataset,
 )
 
 # from happypose.toolbox.datasets.urdf_dataset import UrdfDataset
-
 # TODO: move urdf utilities to happypose toolbox
 from happypose.pose_estimators.cosypose.cosypose.libmesh.urdf_utils import (
     convert_rigid_body_dataset_to_urdfs,
@@ -39,7 +36,9 @@ def __init__(
         elif isinstance(asset_dataset, RigidObjectDataset):
             # Build urdfs files from RigidObjectDataset
             urdf_dir = LOCAL_DATA_DIR / "urdfs" / "tmp"
-            convert_rigid_body_dataset_to_urdfs(asset_dataset, urdf_dir, objname2label_file_name="objname2label.json")
+            convert_rigid_body_dataset_to_urdfs(
+                asset_dataset, urdf_dir, objname2label_file_name="objname2label.json"
+            )
             self.urdf_ds = UrdfDataset(urdf_dir, "objname2label.json")
             # BodyCache assumes unique scale for all objects (true for bop datasets)
             self.urdf_ds.index["scale"] = asset_dataset[0].scale