Merge pull request #139 from agimus-project/132-example-data-and-scri…

…pts-not-unified

132 example data and scripts not unified

.github/workflows/conda-test.yml

@@ -60,11 +60,11 @@ jobs:
           python -m happypose.toolbox.utils.download \
             --megapose_models \
             --examples \
-              crackers_example \
+              barbecue-sauce \
             --cosypose_models  \
-              detector-bop-ycbv-pbr--970850 \
-              coarse-bop-ycbv-pbr--724183 \
-              refiner-bop-ycbv-pbr--604090
+              detector-bop-hope-pbr--15246 \
+              coarse-bop-hope-pbr--225203 \
+              refiner-bop-hope-pbr--955392
 
       - name: Run tests
         run: python -m unittest

.github/workflows/pip-test.yml

@@ -38,11 +38,11 @@ jobs:
           python -m happypose.toolbox.utils.download \
             --megapose_models \
             --examples \
-              crackers_example \
+              barbecue-sauce \
             --cosypose_models  \
-              detector-bop-ycbv-pbr--970850 \
-              coarse-bop-ycbv-pbr--724183 \
-              refiner-bop-ycbv-pbr--604090
+              detector-bop-hope-pbr--15246 \
+              coarse-bop-hope-pbr--225203 \
+              refiner-bop-hope-pbr--955392
 
       - name: Run tests
         run: python -m unittest

.github/workflows/poetry-test.yml

@@ -39,11 +39,11 @@ jobs:
           poetry run python -m happypose.toolbox.utils.download \
             --megapose_models \
             --examples \
-              crackers_example \
+              barbecue-sauce \
             --cosypose_models  \
-              detector-bop-ycbv-pbr--970850 \
-              coarse-bop-ycbv-pbr--724183 \
-              refiner-bop-ycbv-pbr--604090
+              detector-bop-hope-pbr--15246 \
+              coarse-bop-hope-pbr--225203 \
+              refiner-bop-hope-pbr--955392
 
       - name: Run tests
         run: poetry run coverage run --source=happypose -m unittest

docs/book/cosypose/download_data.md

@@ -47,17 +47,23 @@ Notes:
 ## Models for minimal version
 
 ```sh
- #ycbv
-  python -m happypose.toolbox.utils.download --cosypose_models \
-            detector-bop-ycbv-pbr--970850 \
-            coarse-bop-ycbv-pbr--724183 \
-            refiner-bop-ycbv-pbr--604090
-
- #tless
-  python -m happypose.toolbox.utils.download --cosypose_models \
-            detector-bop-tless-pbr--873074 \
-            coarse-bop-tless-pbr--506801 \
-            refiner-bop-tless-pbr--233420
+# hope
+python -m happypose.toolbox.utils.download --cosypose_models \
+          detector-bop-hope-pbr--15246 \
+          coarse-bop-hope-pbr--225203 \
+          refiner-bop-hope-pbr--955392
+
+# ycbv
+python -m happypose.toolbox.utils.download --cosypose_models \
+          detector-bop-ycbv-pbr--970850 \
+          coarse-bop-ycbv-pbr--724183 \
+          refiner-bop-ycbv-pbr--604090
+
+# tless
+python -m happypose.toolbox.utils.download --cosypose_models \
+          detector-bop-tless-pbr--873074 \
+          coarse-bop-tless-pbr--506801 \
+          refiner-bop-tless-pbr--233420
 ```
 
 ## Pre-trained models for single-view estimator

docs/book/cosypose/inference.md

@@ -5,11 +5,11 @@ Here are provided the minimal commands you have to run in order to run the infer
  ## 1. Download pre-trained pose estimation models
 
 ```sh
- #ycbv
+#hope dataset detector
 python -m happypose.toolbox.utils.download --cosypose_models \
-            detector-bop-ycbv-pbr--970850 \
-            coarse-bop-ycbv-pbr--724183 \
-            refiner-bop-ycbv-pbr--604090
+            detector-bop-hope-pbr--15246 \
+            coarse-bop-hope-pbr--225203 \
+            refiner-bop-hope-pbr--955392
 ```
 
 ## 2. Download YCB-V Dataset
@@ -21,19 +21,17 @@ python -m happypose.toolbox.utils.download --bop_dataset=ycbv
 ## 3. Download the example
 
 ```sh
-cd $HAPPYPOSE_DATA_DIR
-wget https://memmo-data.laas.fr/static/examples.tar.xz
-tar xf examples.tar.xz
+python -m happypose.toolbox.utils.download --examples barbecue-sauce
 ```
 
 ## 4. Run the script
-
+The example contains default outputs for detection and pose prediction
 ```sh
-python -m happypose.pose_estimators.cosypose.cosypose.scripts.run_inference_on_example crackers --run-inference
+python -m happypose.pose_estimators.cosypose.cosypose.scripts.run_inference_on_example barbecue-sauce --run-inference --run-detections --vis-detections --vis-poses
 ```
 
 ## 5. Results
 
-The results are stored in the visualization folder created in the crackers example directory.
+The results are stored in the visualization folder created in the example directory.
 
 ![Inference results](./images/all_results.png)

docs/book/megapose/download_data.md

@@ -17,14 +17,26 @@ python -m happypose.toolbox.utils.download --megapose_models
 # Download pre-trained detection models
 Megapose can use pretrained detectors from CosyPose, which can be downloaded to `$HAPPYPOSE_DATA_DIR/experiments`:
 
-python -m happypose.toolbox.utils.download --cosypose_model detector-bop-hb-pbr--497808
-python -m happypose.toolbox.utils.download --cosypose_model detector-bop-hope-pbr--15246
-python -m happypose.toolbox.utils.download --cosypose_model detector-bop-icbin-pbr--947409
-python -m happypose.toolbox.utils.download --cosypose_model detector-bop-itodd-pbr--509908
-python -m happypose.toolbox.utils.download --cosypose_model detector-bop-lmo-pbr--517542
-python -m happypose.toolbox.utils.download --cosypose_model detector-bop-tless-pbr--873074
-python -m happypose.toolbox.utils.download --cosypose_model detector-bop-tudl-pbr--728047
-python -m happypose.toolbox.utils.download --cosypose_model detector-bop-ycbv-pbr--970850
+```sh
+# hope
+python -m happypose.toolbox.utils.download --cosypose_models \
+          detector-bop-hope-pbr--15246 \
+          coarse-bop-hope-pbr--225203 \
+          refiner-bop-hope-pbr--955392
+
+# ycbv
+
+python -m happypose.toolbox.utils.download --cosypose_models \
+          detector-bop-ycbv-pbr--970850 \
+          coarse-bop-ycbv-pbr--724183 \
+          refiner-bop-ycbv-pbr--604090
+
+# tless
+python -m happypose.toolbox.utils.download --cosypose_models \
+          detector-bop-tless-pbr--873074 \
+          coarse-bop-tless-pbr--506801 \
+          refiner-bop-tless-pbr--233420
+```
 
 # Dataset
 

docs/book/megapose/evaluate.md

@@ -119,10 +119,12 @@ module load module load anaconda-py3/2023.03
 conda activate happypose
 
 cd happypose
-# python -m happypose.pose_estimators.megapose.src.megapose.scripts.run_inference_on_example barbecue-sauce --run-inference --vis-outputs
-python -m happypose.pose_estimators.cosypose.cosypose.scripts.run_inference_on_example crackers --run-inference
+# Assuming you have downloaded the example and models
+python -m happypose.pose_estimators.cosypose.cosypose.scripts.run_inference_on_example barbecue-sauce --run-inference
+# python -m happypose.pose_estimators.megapose.scripts.run_inference_on_example barbecue-sauce --run-inference
 ```
 
+
 ```bash
 # evaluation.slurm
 #!/bin/bash

docs/book/megapose/inference.md

@@ -13,9 +13,7 @@ python -m happypose.toolbox.utils.download --megapose_models
 We estimate the pose for a barbecue sauce bottle (from the [HOPE](https://github.com/swtyree/hope-dataset) dataset, not used during training of MegaPose).
 
 ```sh
-cd $HAPPYPOSE_DATA_DIR
-wget https://memmo-data.laas.fr/static/examples.tar.xz
-tar xf examples.tar.xz
+python -m happypose.toolbox.utils.download --examples barbecue-sauce
 ```
 
 The input files are the following:
@@ -25,8 +23,8 @@ $HAPPYPOSE_DATA_DIR/examples/barbecue-sauce/
     image_depth.png
     camera_data.json
     inputs/object_data.json
-    meshes/barbecue-sauce/hope_000002.ply
-    meshes/barbecue-sauce/hope_000002.png
+    meshes/hope-obj_000002.ply
+    meshes/hope-obj_000002.png
 ```
 - `image_rgb.png` is a RGB image of the scene. We recommend using a 4:3 aspect ratio.
 - `image_depth.png` (optional) contains depth measurements, with values in `mm`. You can leave out this file if you don't have depth measurements.
@@ -36,9 +34,9 @@ $HAPPYPOSE_DATA_DIR/examples/barbecue-sauce/
 
 - `inputs/object_data.json` contains a list of object detections. For each detection, the 2D bounding box in the image  (in `[xmin, ymin, xmax, ymax]` format), and the label of the object are provided. In this example, there is a single object detection. The bounding box is only used for computing an initial depth estimate of the object which is then refined by our approach. The bounding box does not need to be extremly precise (see below).
 
-    `[{"label": "barbecue-sauce", "bbox_modal": [384, 234, 522, 455]}]`
+    `[{"label": "hope-obj_000002", "bbox_modal": [384, 234, 522, 455]}]`
 
-- `meshes/barbecue-sauce` is a directory containing the object's mesh. Mesh units are expected to be in millimeters. In this example, we use a mesh in `.ply` format. The code also supports `.obj` meshes but you will have to make sure that the objects are rendered correctly with our renderer.
+- `meshes` is a directory containing the object's mesh. Mesh units are expected to be in millimeters. In this example, we use a mesh in `.ply` format. The code also supports `.obj` meshes but you will have to make sure that the objects are rendered correctly with our renderer.
 
 
 You can visualize input detections using :
@@ -52,27 +50,23 @@ python -m happypose.pose_estimators.megapose.scripts.run_inference_on_example ba
 ## 3. Run pose estimation and visualize results
 Run inference with the following command:
 ```sh
-python -m happypose.pose_estimators.megapose.scripts.run_inference_on_example barbecue-sauce --run-inference
+python -m happypose.pose_estimators.megapose.scripts.run_inference_on_example barbecue-sauce --run-inference --vis-poses
 ```
 by default, the model only uses the RGB input. You can use of our RGB-D megapose models using the `--model` argument. Please see our [Model Zoo](#model-zoo) for all models available.
 
 The previous command will generate the following file:
 
 ```sh
 $HAPPYPOSE_DATA_DIR/examples/barbecue-sauce/
-    outputs/object_data.json
+    outputs/object_data_inf.json
 ```
+A default `object_data.json` is provided if you prefer not to run the model.
 
 This file contains a list of objects with their estimated poses . For each object, the estimated pose is noted `TWO` (the world coordinate frame correspond to the camera frame). It is composed of a quaternion and the 3D translation:
 
     [{"label": "barbecue-sauce", "TWO": [[0.5453961536730983, 0.6226545207599095, -0.43295293693197473, 0.35692612413663855], [0.10723329335451126, 0.07313819974660873, 0.45735278725624084]]}]
 
-Finally, you can visualize the results using:
-
-```sh
-python -m happypose.pose_estimators.megapose.scripts.run_inference_on_example barbecue-sauce --run-inference --vis-outputs
-```
-which write several visualization files:
+The `--vis-poses` options write several visualization files:
 
 ```sh
 $HAPPYPOSE_DATA_DIR/examples/barbecue-sauce/

happypose/pose_estimators/cosypose/cosypose/datasets/bop_object_datasets.py

@@ -1,19 +1,22 @@
 import json
 from pathlib import Path
+from typing import Union
 
 
 class BOPObjectDataset:
-    def __init__(self, ds_dir):
+    def __init__(self, ds_dir, label_format: Union[None, str] = None):
         ds_dir = Path(ds_dir)
         infos_file = ds_dir / "models_info.json"
         infos = json.loads(infos_file.read_text())
         objects = []
         for obj_id, bop_info in infos.items():
             obj_id = int(obj_id)
-            obj_label = f"obj_{obj_id:06d}"
-            mesh_path = (ds_dir / obj_label).with_suffix(".ply").as_posix()
+            label = f"obj_{obj_id:06d}"
+            mesh_path = (ds_dir / label).with_suffix(".ply").as_posix()
+            if label_format is not None:
+                label = label_format.format(label=label)
             obj = {
-                "label": obj_label,
+                "label": label,
                 "category": None,
                 "mesh_path": mesh_path,
                 "mesh_units": "mm",

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

@@ -32,8 +32,7 @@
     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,
+    load_model_cosypose,
 )
 from happypose.pose_estimators.megapose.evaluation.eval_config import EvalConfig
 from happypose.pose_estimators.megapose.evaluation.evaluation_runner import (
@@ -85,46 +84,22 @@ def load_pose_models(coarse_run_id, refiner_run_id, n_workers):
     # 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 = make_object_dataset("ycbv")
-    mesh_db = MeshDataBase.from_object_ds(object_dataset)
     renderer = Panda3dBatchRenderer(
         object_dataset,
         n_workers=n_workers,
         preload_cache=False,
     )
+    mesh_db = MeshDataBase.from_object_ds(object_dataset)
     mesh_db_batched = mesh_db.batched().to(device)
 
-    def load_model(run_id):
-        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_pose(cfg)
-        if cfg.train_refiner:
-            model = create_model_refiner(
-                cfg,
-                renderer=renderer,
-                mesh_db=mesh_db_batched,
-            )
-        else:
-            model = create_model_coarse(cfg, renderer=renderer, mesh_db=mesh_db_batched)
-        ckpt = torch.load(run_dir / "checkpoint.pth.tar", map_location=device)
-        ckpt = ckpt["state_dict"]
-        model.load_state_dict(ckpt)
-        model = model.to(device).eval()
-        model.cfg = cfg
-        model.config = cfg
-        return model
-
-    coarse_model = load_model(coarse_run_id)
-    refiner_model = load_model(refiner_run_id)
+    coarse_model = load_model_cosypose(
+        EXP_DIR / coarse_run_id, renderer, mesh_db_batched, device
+    )
+    refiner_model = load_model_cosypose(
+        EXP_DIR / refiner_run_id, renderer, mesh_db_batched, device
+    )
     return coarse_model, refiner_model, mesh_db
 
 

happypose/pose_estimators/cosypose/cosypose/integrated/pose_estimator.py

@@ -1,6 +1,6 @@
 import time
 from collections import defaultdict
-from typing import Any, Optional, Tuple
+from typing import Any, List, Optional, Tuple
 
 import numpy as np
 import torch
@@ -20,6 +20,7 @@
     ObservationTensor,
     PoseEstimatesType,
 )
+from happypose.toolbox.inference.utils import filter_detections
 from happypose.toolbox.utils.tensor_collection import PandasTensorCollection
 
 logger = get_logger(__name__)
@@ -146,6 +147,7 @@ def run_inference_pipeline(
         coarse_estimates: Optional[PoseEstimatesType] = None,
         detection_th: float = 0.7,
         mask_th: float = 0.8,
+        labels_to_keep: List[str] = None,
     ) -> Tuple[PoseEstimatesType, dict]:
         timing_str = ""
         timer = SimpleTimer()
@@ -180,6 +182,10 @@ def run_inference_pipeline(
             assert detections is not None
             assert self.coarse_model is not None
             assert n_coarse_iterations > 0
+
+            if labels_to_keep is not None:
+                detections = filter_detections(detections, labels_to_keep)
+
             K = observation.K
             data_TCO_init = self.make_TCO_init(detections, K)
             coarse_preds, coarse_extra_data = self.forward_coarse_model(

happypose/pose_estimators/cosypose/cosypose/integrated/pose_predictor.py

@@ -18,6 +18,7 @@
 
 
 class CoarseRefinePosePredictor(PoseEstimationModule):
+    # TODO: deprecate in favor of PosePredictor
     def __init__(self, coarse_model=None, refiner_model=None, bsz_objects=64):
         super().__init__()
         self.coarse_model = coarse_model

happypose/pose_estimators/cosypose/cosypose/lib3d/rigid_mesh_database.py

@@ -92,7 +92,11 @@ def n_sym_mapping(self):
         return {label: obj["n_sym"] for label, obj in self.infos.items()}
 
     def select(self, labels):
-        ids = [self.label_to_id[label] for label in labels]
+        try:
+            ids = [self.label_to_id[label] for label in labels]
+        except KeyError as e:
+            print("self.label_to_id.keys(): ", list(self.label_to_id.keys()))
+            raise e
         return Meshes(
             infos=[self.infos[label] for label in labels],
             labels=self.labels[ids],