Fix some cosypose evaluation bugs

docs/book/cosypose/evaluate.md

@@ -1,11 +1,17 @@
-# Evaluating CosyPose
+# CosyPose single view evaluation
 
 Please make sure you followed the steps relative to the evaluation in the main readme.
 
-Please run the following command to evaluate on YCBV dataset
+To evaluate on YCBV dataset:
 
 ```
-python -m happypose.pose_estimators.cosypose.cosypose.scripts.run_full_cosypose_eval_new detector_run_id=bop_pbr coarse_run_id=coarse-bop-ycbv-pbr--724183 refiner_run_id=refiner-bop-ycbv-pbr--604090 ds_names=["ycbv.bop19"] result_id=ycbv-debug detection_coarse_types=[["detector","S03_grid"]]
+python -m happypose.pose_estimators.cosypose.cosypose.scripts.run_full_cosypose_eval_new detector_run_id=bop_pbr coarse_run_id=coarse-bop-ycbv-pbr--724183 refiner_run_id=refiner-bop-ycbv-pbr--604090 ds_names=["ycbv.bop19"] result_id=ycbv-debug detection_coarse_types=["detector"] inference.renderer=bullet inference.n_workers=0
 ```
 
-The other BOP datasets are supported as long as you download the correspond models.
+To change the renderer from bullet (originally used by cosypose) to panda3d:
+```
+python -m happypose.pose_estimators.cosypose.cosypose.scripts.run_full_cosypose_eval_new detector_run_id=bop_pbr coarse_run_id=coarse-bop-ycbv-pbr--724183 refiner_run_id=refiner-bop-ycbv-pbr--604090 ds_names=["ycbv.bop19"] result_id=ycbv-debug detection_coarse_types=["detector"] inference.renderer=panda3d inference.n_workers=1
+```
+
+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"].

happypose/pose_estimators/cosypose/cosypose/config.py

@@ -17,9 +17,13 @@
 PROJECT_ROOT = Path(happypose.__file__).parent.parent
 PROJECT_DIR = PROJECT_ROOT
 DATA_DIR = PROJECT_DIR / "data"
+# user should set "HAPPYPOSE_DATA_DIR" env or create project dir "local_data"
 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'?"
 TEST_DATA_DIR = LOCAL_DATA_DIR
 DASK_LOGS_DIR = LOCAL_DATA_DIR / "dasklogs"
 SYNT_DS_DIR = LOCAL_DATA_DIR / "synt_datasets"
@@ -36,7 +40,6 @@
 DEPS_DIR = PROJECT_DIR / "deps"
 CACHE_DIR = LOCAL_DATA_DIR / "joblib_cache"
 
-assert LOCAL_DATA_DIR.exists()
 CACHE_DIR.mkdir(exist_ok=True)
 TEST_DATA_DIR.mkdir(exist_ok=True)
 DASK_LOGS_DIR.mkdir(exist_ok=True)

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

@@ -5,23 +5,20 @@
 
 
 class UrdfDataset:
-    def __init__(self, urdf_ds_dir, label_filename="objname2label.json"):
+    def __init__(self, urdf_ds_dir, objname2label_file_name="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
+        with open(urdf_ds_dir / objname2label_file_name, "r") as f:
+            objname2label = json.load(f)
         index = []
         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]
+                obj_name = obj_dir.name
                 if objname2label is None:
-                    label = obj_dir.name
+                    label = obj_name
                 else:
-                    label = objname2label[obj_dir.name]
+                    label = objname2label[obj_name]
                 infos = {
                     "label": label,
                     "urdf_path": urdf_path.as_posix(),

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

@@ -8,10 +8,6 @@
 from omegaconf import OmegaConf
 
 # MegaPose
-import happypose
-import happypose.pose_estimators.megapose.evaluation.evaluation_runner
-import happypose.toolbox.datasets.datasets_cfg
-import happypose.toolbox.inference.utils
 from happypose.pose_estimators.cosypose.cosypose.config import EXP_DIR
 from happypose.pose_estimators.cosypose.cosypose.evaluation.prediction_runner import (
     PredictionRunner,
@@ -20,14 +16,14 @@
 from happypose.pose_estimators.cosypose.cosypose.integrated.pose_estimator import (
     PoseEstimator,
 )
-
-# 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,
 )
+
+# Detection
 from happypose.pose_estimators.cosypose.cosypose.training.pose_models_cfg import (
     check_update_config as check_update_config_pose,
 )
@@ -44,23 +40,22 @@
 from happypose.pose_estimators.megapose.evaluation.runner_utils import format_results
 
 # Pose estimator
-from happypose.toolbox.datasets.datasets_cfg import make_object_dataset
+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.renderer.panda3d_batch_renderer import Panda3dBatchRenderer
 from happypose.toolbox.utils.distributed import get_rank, get_tmp_dir
 from happypose.toolbox.utils.logging import get_logger
 
-# """" Temporary imports
-
-
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 logger = get_logger(__name__)
 
 
 def load_detector(run_id, ds_name):
     run_dir = EXP_DIR / 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_detector(cfg)
     label_to_category_id = cfg.label_to_category_id
@@ -75,18 +70,34 @@ def load_detector(run_id, ds_name):
     return model
 
 
-def load_pose_models(coarse_run_id, refiner_run_id, n_workers):
+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.FullLoader)
     cfg = yaml.load((run_dir / "config.yaml").read_text(), Loader=yaml.UnsafeLoader)
     cfg = check_update_config_pose(cfg)
 
-    object_dataset = make_object_dataset("ycbv")
-    renderer = Panda3dBatchRenderer(
-        object_dataset,
-        n_workers=n_workers,
-        preload_cache=False,
-    )
+    if renderer_type == "panda3d":
+        from happypose.toolbox.renderer.panda3d_batch_renderer import (
+            Panda3dBatchRenderer,
+        )
+
+        renderer = Panda3dBatchRenderer(
+            object_dataset,
+            n_workers=n_workers,
+            preload_cache=True,
+        )
+    elif renderer_type == "bullet":
+        from happypose.toolbox.renderer.bullet_batch_renderer import BulletBatchRenderer
+
+        renderer = BulletBatchRenderer(
+            object_dataset,
+            n_workers=n_workers,
+            preload_cache=True,
+        )
+    else:
+        raise ValueError(f"Renderer {renderer_type} not supported")
+
     mesh_db = MeshDataBase.from_object_ds(object_dataset)
     mesh_db_batched = mesh_db.batched().to(device)
 
@@ -153,14 +164,16 @@ def run_eval(
     cfg.save_dir = str(save_dir)
 
     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]
 
     # Load the dataset
     ds_kwargs = {"load_depth": False}
-    scene_ds = happypose.toolbox.datasets.datasets_cfg.make_scene_dataset(
+    scene_ds = make_scene_dataset(
         cfg.ds_name,
         **ds_kwargs,
     )
-    urdf_ds_name, obj_ds_name = happypose.toolbox.datasets.datasets_cfg.get_obj_ds_info(
+    urdf_ds_name, obj_ds_name = get_obj_ds_info(
         cfg.ds_name,
     )
 
@@ -173,7 +186,7 @@ def run_eval(
     # Load detector model
     if cfg.inference.detection_type == "detector":
         assert cfg.detector_run_id is not None
-        detector_model = load_detector(cfg.detector_run_id, cfg.ds_name)
+        detector_model = load_detector(cfg.detector_run_id, ds_name_short)
     elif cfg.inference.detection_type == "gt":
         detector_model = None
     else:
@@ -186,12 +199,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(obj_ds_name)
+    object_ds = make_object_dataset(ds_name_short)
 
-    coarse_model, refiner_model, mesh_db = load_pose_models(
+    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=8,
+        n_workers=cfg.inference.n_workers,
+        renderer_type=cfg.inference.renderer,
     )
 
     renderer = refiner_model.renderer

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

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

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

@@ -16,12 +16,12 @@ def convert_rigid_body_dataset_to_urdfs(
     urdf_dir: Path,
     texture_size=(1024, 1024),
     override=True,
-    label2objname_file_name: str = "objname2label.json",
+    objname2label_file_name: str = "objname2label.json",
 ):
     """
     Converts a RigidObjectDataset into a directory of urdf files with structure:
 
-    urdf_dir/<label2objname_file_name>.json
+    urdf_dir/<objname2label_file_name>.json
     urdf_dir/obj_000001/obj_000001.mtl
                         obj_000001.obj
                         obj_000001_texture.png
@@ -31,8 +31,20 @@ def convert_rigid_body_dataset_to_urdfs(
                         obj_000002_texture.png
                         obj_000002.urdf
 
-    <label2objname_file_name>.json: stores a map between object file names (e.g. obj_000002) and
+    <objname2label_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")
+
+    Args:
+    ----
+        rb_ds (RigidObjectDataset):
+        urdf_dir (Path):
+        texture_size (tuple):
+        override (bool): override directory
+        objname2label_file_name (str):
+
+    Returns:
+    -------
+        List[SceneData]: _description_
     """
     if override and urdf_dir.exists():
         shutil.rmtree(urdf_dir, ignore_errors=True)
@@ -42,6 +54,7 @@ def convert_rigid_body_dataset_to_urdfs(
     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
+        # import pdb; pdb.set_trace()
         # Create object folder
         obj_urdf_dir = urdf_dir / objname
         obj_urdf_dir.mkdir(exist_ok=True)  # urdf_dir/obj_000002/ created
@@ -55,7 +68,7 @@ def convert_rigid_body_dataset_to_urdfs(
         urdf_path = obj_path.with_suffix(".urdf")
         obj_to_urdf(obj_path, urdf_path)
 
-    with open(urdf_dir / label2objname_file_name, "w") as fp:
+    with open(urdf_dir / objname2label_file_name, "w") as fp:
         json.dump(objname2label, fp)
 
 

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

@@ -10,10 +10,12 @@
 from happypose.pose_estimators.cosypose.cosypose.evaluation.evaluation import run_eval
 
 # MegaPose
-from happypose.pose_estimators.megapose.bop_config import PBR_DETECTORS
+from happypose.pose_estimators.megapose.bop_config import (
+    PBR_DETECTORS,
+    SYNT_REAL_DETECTORS,
+)
 from happypose.pose_estimators.megapose.config import (
     DEBUG_RESULTS_DIR,
-    MODELNET_TEST_CATEGORIES,
     RESULTS_DIR,
 )
 from happypose.pose_estimators.megapose.evaluation.bop import run_evaluation
@@ -36,18 +38,6 @@
 
 logger = get_logger(__name__)
 
-BOP_DATASET_NAMES = [
-    "lm",
-    "lmo",
-    "tless",
-    "tudl",
-    "icbin",
-    "itodd",
-    "hb",
-    "ycbv",
-    # 'hope',
-]
-
 BOP_TEST_DATASETS = [
     "lmo.bop19",
     "tless.bop19",
@@ -59,21 +49,14 @@
 ]
 
 
-MODELNET_TEST_DATASETS = [
-    f"modelnet.{category}.test" for category in MODELNET_TEST_CATEGORIES
-]
-
-
 def create_eval_cfg(
     cfg: EvalConfig,
     detection_type: str,
-    coarse_estimation_type: str,
     ds_name: str,
 ) -> Tuple[str, EvalConfig]:
     cfg = copy.deepcopy(cfg)
 
     cfg.inference.detection_type = detection_type
-    cfg.inference.coarse_estimation_type = coarse_estimation_type
     cfg.ds_name = ds_name
 
     if detection_type == "detector":
@@ -82,13 +65,15 @@ def create_eval_cfg(
         ds_name_root = cfg.ds_name.split(".")[0]
         if cfg.detector_run_id == "bop_pbr":
             cfg.detector_run_id = PBR_DETECTORS[ds_name_root]
+        elif cfg.detector_run_id == "bop_synt_real":
+            cfg.detector_run_id = SYNT_REAL_DETECTORS[ds_name_root]
     elif detection_type == "gt":
         pass
     else:
-        msg = f"Unknown detector type {cfg.detector_type}"
+        msg = f"Unknown detector type {detection_type}"
         raise ValueError(msg)
 
-    name = generate_save_key(detection_type, coarse_estimation_type)
+    name = generate_save_key(detection_type, "cosycoarse")
 
     return name, cfg
 
@@ -109,11 +94,10 @@ def run_full_eval(cfg: FullEvalConfig) -> None:
     for ds_name in cfg.ds_names:
         # create the EvalConfig objects that we will call `run_eval` on
         eval_configs: Dict[str, EvalConfig] = {}
-        for detection_type, coarse_estimation_type in cfg.detection_coarse_types:
+        for detection_type in cfg.detection_coarse_types:
             name, cfg_ = create_eval_cfg(
                 cfg,
                 detection_type,
-                coarse_estimation_type,
                 ds_name,
             )
             eval_configs[name] = cfg_