diff --git a/.github/workflows/conda-test.yml b/.github/workflows/conda-test.yml index 2fe955e4..27fdba6e 100644 --- a/.github/workflows/conda-test.yml +++ b/.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 diff --git a/.github/workflows/pip-test.yml b/.github/workflows/pip-test.yml index 97af3ee5..7e05aac7 100644 --- a/.github/workflows/pip-test.yml +++ b/.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 diff --git a/.github/workflows/poetry-test.yml b/.github/workflows/poetry-test.yml index c3c95fe9..2d177472 100644 --- a/.github/workflows/poetry-test.yml +++ b/.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 diff --git a/docs/book/cosypose/download_data.md b/docs/book/cosypose/download_data.md index faefa55e..085e3584 100644 --- a/docs/book/cosypose/download_data.md +++ b/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 diff --git a/docs/book/cosypose/images/all_results.png b/docs/book/cosypose/images/all_results.png index ac3e88f1..9f28eddc 100644 Binary files a/docs/book/cosypose/images/all_results.png and b/docs/book/cosypose/images/all_results.png differ diff --git a/docs/book/cosypose/inference.md b/docs/book/cosypose/inference.md index 90e50c6e..3a0e3b41 100644 --- a/docs/book/cosypose/inference.md +++ b/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) diff --git a/docs/book/megapose/download_data.md b/docs/book/megapose/download_data.md index 562a278d..552a5f80 100644 --- a/docs/book/megapose/download_data.md +++ b/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 diff --git a/docs/book/megapose/evaluate.md b/docs/book/megapose/evaluate.md index c72ffac6..c1adf1ec 100644 --- a/docs/book/megapose/evaluate.md +++ b/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 diff --git a/docs/book/megapose/inference.md b/docs/book/megapose/inference.md index b44c0d46..7c08c792 100644 --- a/docs/book/megapose/inference.md +++ b/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,7 +50,7 @@ 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. @@ -60,19 +58,15 @@ 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/ diff --git a/happypose/pose_estimators/cosypose/cosypose/datasets/bop_object_datasets.py b/happypose/pose_estimators/cosypose/cosypose/datasets/bop_object_datasets.py index 22fa5301..7393c05d 100644 --- a/happypose/pose_estimators/cosypose/cosypose/datasets/bop_object_datasets.py +++ b/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", diff --git a/happypose/pose_estimators/cosypose/cosypose/evaluation/evaluation.py b/happypose/pose_estimators/cosypose/cosypose/evaluation/evaluation.py index 4a22cf77..e32ddb25 100644 --- a/happypose/pose_estimators/cosypose/cosypose/evaluation/evaluation.py +++ b/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 diff --git a/happypose/pose_estimators/cosypose/cosypose/integrated/pose_estimator.py b/happypose/pose_estimators/cosypose/cosypose/integrated/pose_estimator.py index 2e23eb70..34d9bda0 100644 --- a/happypose/pose_estimators/cosypose/cosypose/integrated/pose_estimator.py +++ b/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( diff --git a/happypose/pose_estimators/cosypose/cosypose/integrated/pose_predictor.py b/happypose/pose_estimators/cosypose/cosypose/integrated/pose_predictor.py index ae4291d7..57451618 100644 --- a/happypose/pose_estimators/cosypose/cosypose/integrated/pose_predictor.py +++ b/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 diff --git a/happypose/pose_estimators/cosypose/cosypose/lib3d/rigid_mesh_database.py b/happypose/pose_estimators/cosypose/cosypose/lib3d/rigid_mesh_database.py index db8ac1b3..68881f1e 100644 --- a/happypose/pose_estimators/cosypose/cosypose/lib3d/rigid_mesh_database.py +++ b/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], diff --git a/happypose/pose_estimators/cosypose/cosypose/libmesh/urdf_utils.py b/happypose/pose_estimators/cosypose/cosypose/libmesh/urdf_utils.py index fce477f8..668d4a20 100644 --- a/happypose/pose_estimators/cosypose/cosypose/libmesh/urdf_utils.py +++ b/happypose/pose_estimators/cosypose/cosypose/libmesh/urdf_utils.py @@ -30,15 +30,17 @@ def convert_rigid_body_dataset_to_urdfs( obj_to_urdf(obj_path, urdf_path) -def ply_to_obj(ply_path, obj_path, texture_size=None): +def ply_to_obj(ply_path: Path, obj_path: Path, texture_size=None): + assert obj_path.suffix == ".obj" mesh = trimesh.load(ply_path) 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" + if mesh.visual.defined: + 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"} diff --git a/happypose/pose_estimators/cosypose/cosypose/rendering/__init__.py b/happypose/pose_estimators/cosypose/cosypose/rendering/__init__.py deleted file mode 100644 index e69de29b..00000000 diff --git a/happypose/pose_estimators/cosypose/cosypose/scripts/convert_models_to_urdf.py b/happypose/pose_estimators/cosypose/cosypose/scripts/convert_models_to_urdf.py index 1d65c77b..00569b26 100644 --- a/happypose/pose_estimators/cosypose/cosypose/scripts/convert_models_to_urdf.py +++ b/happypose/pose_estimators/cosypose/cosypose/scripts/convert_models_to_urdf.py @@ -3,26 +3,37 @@ from tqdm import tqdm -from happypose.pose_estimators.cosypose.cosypose.datasets.datasets_cfg import ( - make_object_dataset, -) from happypose.pose_estimators.cosypose.cosypose.libmesh.urdf_utils import ( obj_to_urdf, ply_to_obj, ) +# from happypose.pose_estimators.cosypose.cosypose.datasets.datasets_cfg import \ +# make_object_dataset +from happypose.toolbox.datasets.datasets_cfg import make_object_dataset + def convert_bop_dataset_to_urdfs( obj_ds_name: str, urdf_dir: Path, texture_size=(1024, 1024) ): + """ + For each object, generate these files: + + {path_to_urdf_dir}/{ds_name}/{obj_label}/{obj_label}.obj + {path_to_urdf_dir}/{ds_name}/{obj_label}/{obj_label}.obj.mtl + {path_to_urdf_dir}/{ds_name}/{obj_label}/{obj_label}.png + {path_to_urdf_dir}/{ds_name}/{obj_label}/{obj_label}.urdf + """ obj_dataset = make_object_dataset(obj_ds_name) - urdf_dir.mkdir(exist_ok=True, parents=True) + urdf_ds_dir = urdf_dir / obj_ds_name + urdf_ds_dir.mkdir(exist_ok=True, parents=True) for n in tqdm(range(len(obj_dataset))): obj = obj_dataset[n] - ply_path = Path(obj["mesh_path"]) - out_dir = urdf_dir / obj["label"] - out_dir.mkdir(exist_ok=True) - obj_path = out_dir / ply_path.with_suffix(".obj").name + ply_path = obj.mesh_path + obj_name = ply_path.with_suffix("").name + obj_urdf_dir = urdf_ds_dir / obj_name + obj_urdf_dir.mkdir(exist_ok=True) + obj_path = (obj_urdf_dir / obj_name).with_suffix(".obj") ply_to_obj(ply_path, obj_path, texture_size=texture_size) obj_to_urdf(obj_path, obj_path.with_suffix(".urdf")) @@ -34,8 +45,8 @@ def main(): from happypose.pose_estimators.cosypose.cosypose.config import LOCAL_DATA_DIR - urdf_dir = LOCAL_DATA_DIR / "urdf" - convert_bop_dataset_to_urdfs(urdf_dir, args.models) + urdf_dir = LOCAL_DATA_DIR / "urdfs" + convert_bop_dataset_to_urdfs(args.models, urdf_dir) if __name__ == "__main__": diff --git a/happypose/pose_estimators/cosypose/cosypose/scripts/run_bop_inference.py b/happypose/pose_estimators/cosypose/cosypose/scripts/run_bop_inference.py index 33525ec4..affa766d 100644 --- a/happypose/pose_estimators/cosypose/cosypose/scripts/run_bop_inference.py +++ b/happypose/pose_estimators/cosypose/cosypose/scripts/run_bop_inference.py @@ -58,8 +58,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.cosypose.cosypose.utils.distributed import ( get_rank, @@ -75,6 +74,8 @@ logger = get_logger(__name__) +device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + def load_detector(run_id): run_dir = EXP_DIR / run_id @@ -102,30 +103,12 @@ def load_pose_models(coarse_run_id, refiner_run_id=None, n_workers=8): renderer = BulletBatchRenderer(object_set=cfg.urdf_ds_name, n_workers=n_workers) mesh_db_batched = mesh_db.batched().cuda() - def load_model(run_id): - if run_id is None: - return None - run_dir = EXP_DIR / run_id - cfg = yaml.load((run_dir / "config.yaml").read_text(), Loader=yaml.FullLoader) - 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") - ckpt = ckpt["state_dict"] - model.load_state_dict(ckpt) - model = model.cuda().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 + ) model = CoarseRefinePosePredictor( coarse_model=coarse_model, refiner_model=refiner_model, diff --git a/happypose/pose_estimators/cosypose/cosypose/scripts/run_cosypose_eval.py b/happypose/pose_estimators/cosypose/cosypose/scripts/run_cosypose_eval.py index 7759c156..73e438c9 100644 --- a/happypose/pose_estimators/cosypose/cosypose/scripts/run_cosypose_eval.py +++ b/happypose/pose_estimators/cosypose/cosypose/scripts/run_cosypose_eval.py @@ -13,7 +13,6 @@ import pandas as pd import torch import torch.multiprocessing -import yaml import happypose.pose_estimators.cosypose.cosypose.utils.tensor_collection as tc from happypose.pose_estimators.cosypose.cosypose.config import ( @@ -54,9 +53,7 @@ MeshDataBase, ) from happypose.pose_estimators.cosypose.cosypose.training.pose_models_cfg import ( - check_update_config, - create_model_coarse, - create_model_refiner, + load_model_cosypose, ) from happypose.pose_estimators.cosypose.cosypose.utils.distributed import ( get_rank, @@ -73,6 +70,8 @@ torch.backends.cudnn.deterministic = True torch.backends.cudnn.benchmark = False +device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + @MEMORY.cache def load_posecnn_results(): @@ -292,30 +291,12 @@ def load_models(coarse_run_id, refiner_run_id=None, n_workers=8, object_set="tle renderer = BulletBatchRenderer(object_set=urdf_ds_name, n_workers=n_workers) mesh_db_batched = mesh_db.batched().cuda() - def load_model(run_id): - if run_id is None: - return None - run_dir = EXP_DIR / run_id - cfg = yaml.load((run_dir / "config.yaml").read_text(), Loader=yaml.FullLoader) - cfg = check_update_config(cfg) - if cfg.train_refiner: - model = create_model_refiner( - cfg, - renderer=renderer, - mesh_db=mesh_db_batched, - ) - ckpt = torch.load(run_dir / "checkpoint.pth.tar") - else: - model = create_model_coarse(cfg, renderer=renderer, mesh_db=mesh_db_batched) - ckpt = torch.load(run_dir / "checkpoint.pth.tar") - ckpt = ckpt["state_dict"] - model.load_state_dict(ckpt) - model = model.cuda().eval() - model.cfg = 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 + ) model = CoarseRefinePosePredictor( coarse_model=coarse_model, refiner_model=refiner_model, diff --git a/happypose/pose_estimators/cosypose/cosypose/scripts/run_inference_on_example.py b/happypose/pose_estimators/cosypose/cosypose/scripts/run_inference_on_example.py index dca59319..16a333b7 100644 --- a/happypose/pose_estimators/cosypose/cosypose/scripts/run_inference_on_example.py +++ b/happypose/pose_estimators/cosypose/cosypose/scripts/run_inference_on_example.py @@ -1,200 +1,118 @@ # Standard Library import argparse import os - -######################## -# Add cosypose to my path -> dirty from pathlib import Path -from typing import Tuple, Union # Third Party -import numpy as np import torch -from bokeh.io import export_png -from bokeh.plotting import gridplot -from PIL import Image +from happypose.pose_estimators.cosypose.cosypose.integrated.pose_estimator import ( + PoseEstimator, +) + +# CosyPose from happypose.pose_estimators.cosypose.cosypose.utils.cosypose_wrapper import ( CosyPoseWrapper, ) -# MegaPose -from happypose.toolbox.datasets.object_dataset import RigidObject, RigidObjectDataset -from happypose.toolbox.datasets.scene_dataset import CameraData, ObjectData -from happypose.toolbox.inference.types import ObservationTensor -from happypose.toolbox.lib3d.transform import Transform - # HappyPose -from happypose.toolbox.renderer import Panda3dLightData -from happypose.toolbox.renderer.panda3d_scene_renderer import Panda3dSceneRenderer -from happypose.toolbox.utils.conversion import convert_scene_observation_to_panda3d +from happypose.toolbox.datasets.object_dataset import RigidObjectDataset +from happypose.toolbox.inference.example_inference_utils import ( + load_detections, + load_object_data, + load_observation_example, + make_detections_visualization, + make_example_object_dataset, + make_poses_visualization, + save_predictions, +) +from happypose.toolbox.inference.types import DetectionsType, ObservationTensor +from happypose.toolbox.inference.utils import filter_detections, load_detector from happypose.toolbox.utils.logging import get_logger, set_logging_level -from happypose.toolbox.visualization.bokeh_plotter import BokehPlotter -from happypose.toolbox.visualization.utils import make_contour_overlay - -######################## - logger = get_logger(__name__) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") -def load_observation( - example_dir: Path, - load_depth: bool = False, -) -> Tuple[np.ndarray, Union[None, np.ndarray], CameraData]: - camera_data = CameraData.from_json((example_dir / "camera_data.json").read_text()) - - rgb = np.array(Image.open(example_dir / "image_rgb.png"), dtype=np.uint8) - assert rgb.shape[:2] == camera_data.resolution - - depth = None - if load_depth: - depth = ( - np.array(Image.open(example_dir / "image_depth.png"), dtype=np.float32) - / 1000 - ) - assert depth.shape[:2] == camera_data.resolution - - return rgb, depth, camera_data - - -def load_observation_tensor( - example_dir: Path, - load_depth: bool = False, -) -> ObservationTensor: - rgb, depth, camera_data = load_observation(example_dir, load_depth) - observation = ObservationTensor.from_numpy(rgb, depth, camera_data.K) - if torch.cuda.is_available(): - observation.cuda() - return observation - - -def make_object_dataset(example_dir: Path) -> RigidObjectDataset: - rigid_objects = [] - mesh_units = "mm" - object_dirs = (example_dir / "meshes").iterdir() - for object_dir in object_dirs: - label = object_dir.name - mesh_path = None - for fn in object_dir.glob("*"): - if fn.suffix in {".obj", ".ply"}: - assert not mesh_path, f"there multiple meshes in the {label} directory" - mesh_path = fn - assert mesh_path, f"couldnt find a obj or ply mesh for {label}" - rigid_objects.append( - RigidObject(label=label, mesh_path=mesh_path, mesh_units=mesh_units), - ) - # TODO: fix mesh units - rigid_object_dataset = RigidObjectDataset(rigid_objects) - return rigid_object_dataset - - -def rendering(predictions, example_dir): - object_dataset = make_object_dataset(example_dir) - - obj_label = object_dataset.list_objects[0].label - pred = predictions.poses[0].numpy() - - # rendering - camera_data = CameraData.from_json((example_dir / "camera_data.json").read_text()) - camera_data.TWC = Transform(np.eye(4)) - renderer = Panda3dSceneRenderer(object_dataset) - # Data necessary for image rendering - object_datas = [ObjectData(label=obj_label, TWO=Transform(pred))] - camera_data, object_datas = convert_scene_observation_to_panda3d( - camera_data, - object_datas, - ) - light_datas = [ - Panda3dLightData( - light_type="ambient", - color=((0.6, 0.6, 0.6, 1)), - ), - ] - renderings = renderer.render_scene( - object_datas, - [camera_data], - light_datas, - render_depth=False, - render_binary_mask=False, - render_normals=False, - copy_arrays=True, - )[0] - return renderings - - -def save_predictions(example_dir, renderings): - rgb_render = renderings.rgb - rgb, _, _ = load_observation(example_dir, load_depth=False) - mask = ~(rgb_render.sum(axis=-1) == 0) - rgb_n_render = rgb.copy() - rgb_n_render[mask] = rgb_render[mask] - - # make the image background a bit fairer than the render - rgb_overlay = np.zeros_like(rgb_render) - rgb_overlay[~mask] = rgb[~mask] * 0.6 + 255 * 0.4 - rgb_overlay[mask] = rgb_render[mask] * 0.8 + 255 * 0.2 - plotter = BokehPlotter() - - fig_rgb = plotter.plot_image(rgb) - - fig_mesh_overlay = plotter.plot_overlay(rgb, renderings.rgb) - contour_overlay = make_contour_overlay( - rgb, - renderings.rgb, - dilate_iterations=1, - color=(0, 255, 0), - )["img"] - fig_contour_overlay = plotter.plot_image(contour_overlay) - fig_all = gridplot( - [[fig_rgb, fig_contour_overlay, fig_mesh_overlay]], - toolbar_location=None, +def setup_pose_estimator(dataset_to_use: str, object_dataset: RigidObjectDataset): + # TODO: remove this wrapper from code base + cosypose = CosyPoseWrapper( + dataset_name=dataset_to_use, object_dataset=object_dataset, n_workers=1 ) - vis_dir = example_dir / "visualizations" - vis_dir.mkdir(exist_ok=True) - export_png(fig_mesh_overlay, filename=vis_dir / "mesh_overlay.png") - export_png(fig_contour_overlay, filename=vis_dir / "contour_overlay.png") - export_png(fig_all, filename=vis_dir / "all_results.png") + + return cosypose.pose_predictor def run_inference( - example_dir: Path, - model_name: str, - dataset_to_use: str, + pose_estimator: PoseEstimator, + observation: ObservationTensor, + detections: DetectionsType, ) -> None: - observation = load_observation_tensor(example_dir) - # TODO: remove this wrapper from code base - CosyPose = CosyPoseWrapper(dataset_name=dataset_to_use, n_workers=8) - predictions = CosyPose.inference(observation) - renderings = rendering(predictions, example_dir) - save_predictions(example_dir, renderings) + observation.to(device) + + data_TCO, extra_data = pose_estimator.run_inference_pipeline( + observation=observation, detections=detections, n_refiner_iterations=3 + ) + print("Timings:") + print(extra_data["timing_str"]) + + return data_TCO.cpu() if __name__ == "__main__": set_logging_level("info") parser = argparse.ArgumentParser() parser.add_argument("example_name") - # parser.add_argument( - # "--model", type=str, default="megapose-1.0-RGB-multi-hypothesis" - # ) - parser.add_argument("--dataset", type=str, default="ycbv") - # parser.add_argument("--vis-detections", action="store_true") - parser.add_argument("--run-inference", action="store_true", default=True) - # parser.add_argument("--vis-outputs", action="store_true") + parser.add_argument("--dataset", type=str, default="hope") + parser.add_argument("--run-detections", action="store_true") + parser.add_argument("--run-inference", action="store_true") + parser.add_argument("--vis-detections", action="store_true") + parser.add_argument("--vis-poses", action="store_true") args = parser.parse_args() data_dir = os.getenv("HAPPYPOSE_DATA_DIR") - assert data_dir + assert data_dir, "Set HAPPYPOSE_DATA_DIR env variable" example_dir = Path(data_dir) / "examples" / args.example_name - dataset_to_use = args.dataset # tless or ycbv - - # if args.vis_detections: - # make_detections_visualization(example_dir) + assert ( + example_dir.exists() + ), "Example {args.example_name} not available, follow download instructions" + dataset_to_use = args.dataset # hope/tless/ycbv + + # Load data + detections = load_detections(example_dir).to(device) + object_dataset = make_example_object_dataset(example_dir) + rgb, depth, camera_data = load_observation_example(example_dir, load_depth=True) + # TODO: cosypose forward does not work if depth is loaded detection + # contrary to megapose + observation = ObservationTensor.from_numpy(rgb, depth=None, K=camera_data.K) + + # Load models + pose_estimator = setup_pose_estimator(args.dataset, object_dataset) + + if args.run_detections: + # TODO: hardcoded detector + detector = load_detector(run_id="detector-bop-hope-pbr--15246", device=device) + # Masks are not used for pose prediction, but are computed by Mask-RCNN anyway + detections = detector.get_detections(observation, output_masks=True) + available_labels = [obj.label for obj in object_dataset.list_objects] + detections = filter_detections(detections, available_labels) + else: + detections = load_detections(example_dir).to(device) if args.run_inference: - run_inference(example_dir, None, dataset_to_use) - - # if args.vis_outputs: - # make_output_visualization(example_dir) + output = run_inference(pose_estimator, observation, detections) + save_predictions(output, example_dir) + + if args.vis_detections: + make_detections_visualization(rgb, detections, example_dir) + + if args.vis_poses: + if args.run_inference: + out_filename = "object_data_inf.json" + else: + out_filename = "object_data.json" + object_datas = load_object_data(example_dir / "outputs" / out_filename) + make_poses_visualization( + rgb, object_dataset, object_datas, camera_data, example_dir + ) diff --git a/happypose/pose_estimators/cosypose/cosypose/simulator/caching.py b/happypose/pose_estimators/cosypose/cosypose/simulator/caching.py index 2b151d38..89265c1f 100644 --- a/happypose/pose_estimators/cosypose/cosypose/simulator/caching.py +++ b/happypose/pose_estimators/cosypose/cosypose/simulator/caching.py @@ -15,8 +15,10 @@ def __init__(self, urdf_ds, client_id): self.away_transform = (0, 0, 1000), (0, 0, 0, 1) def _load_body(self, label): - ds_idx = np.where(self.urdf_ds.index["label"] == label)[0].item() - object_infos = self.urdf_ds[ds_idx].to_dict() + ds_idx = np.where(self.urdf_ds.index["label"] == label)[0] + if len(ds_idx) == 0: + raise ValueError(f'Label {label} not in {self.urdf_ds.index["label"]}') + object_infos = self.urdf_ds[ds_idx.item()].to_dict() body = Body.load( object_infos["urdf_path"], scale=object_infos["scale"], diff --git a/happypose/pose_estimators/cosypose/cosypose/training/evaluation.py b/happypose/pose_estimators/cosypose/cosypose/training/evaluation.py new file mode 100644 index 00000000..4f7d80d8 --- /dev/null +++ b/happypose/pose_estimators/cosypose/cosypose/training/evaluation.py @@ -0,0 +1,323 @@ +# Standard Library +from pathlib import Path +from typing import Any, Dict, Optional + +# Third Party +import torch +import yaml +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, +) +from happypose.pose_estimators.cosypose.cosypose.integrated.detector import Detector +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, +) +from happypose.pose_estimators.cosypose.cosypose.training.pose_models_cfg import ( + check_update_config as check_update_config_pose, +) +from happypose.pose_estimators.cosypose.cosypose.training.pose_models_cfg import ( + load_model_cosypose, +) +from happypose.pose_estimators.megapose.evaluation.eval_config import EvalConfig +from happypose.pose_estimators.megapose.evaluation.evaluation_runner import ( + EvaluationRunner, +) +from happypose.pose_estimators.megapose.evaluation.meters.modelnet_meters import ( + ModelNetErrorMeter, +) +from happypose.pose_estimators.megapose.evaluation.runner_utils import format_results +from happypose.pose_estimators.megapose.inference.icp_refiner import ICPRefiner + +# Pose estimator +from happypose.pose_estimators.megapose.inference.teaserpp_refiner import ( + TeaserppRefiner, +) +from happypose.toolbox.datasets.datasets_cfg import make_object_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 + model = create_model_detector(cfg, len(label_to_category_id)) + 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 + model = Detector(model, ds_name) + return model + + +def load_pose_models(coarse_run_id, refiner_run_id, n_workers): + 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_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") + 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) + + 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 + + +def generate_save_key(detection_type: str, coarse_estimation_type: str) -> str: + return f"{detection_type}+{coarse_estimation_type}" + + +def get_save_dir(cfg: EvalConfig) -> Path: + """Returns a save dir. + + Example: + ------- + .../ycbv.bop19/gt+SO3_grid + + You must remove the '.bop19' from the name in order for the + bop_toolkit_lib to process it correctly. + + """ + save_key = generate_save_key( + cfg.inference.detection_type, + cfg.inference.coarse_estimation_type, + ) + + assert cfg.save_dir is not None + assert cfg.ds_name is not None + save_dir = Path(cfg.save_dir) / cfg.ds_name / save_key + return save_dir + + +def run_eval( + cfg: EvalConfig, + save_dir: Optional[Path] = None, +) -> Dict[str, Any]: + """Run eval for a single setting on a single dataset. + + A single setting is a (detection_type, coarse_estimation_type) such + as ('maskrcnn', 'SO3_grid'). + + Saves the results to the directory below (if one is not passed in). + + cfg.save_dir / ds_name / eval_key / results.pth.tar + + Returns + ------- + dict: If you are rank_0 process, otherwise returns None + + """ + save_key = generate_save_key( + cfg.inference.detection_type, + cfg.inference.coarse_estimation_type, + ) + if save_dir is None: + save_dir = get_save_dir(cfg) + + cfg.save_dir = str(save_dir) + + logger.info(f"Running eval on ds_name={cfg.ds_name} with setting={save_key}") + + # Load the dataset + ds_kwargs = {"load_depth": False} + scene_ds = happypose.toolbox.datasets.datasets_cfg.make_scene_dataset( + cfg.ds_name, + **ds_kwargs, + ) + urdf_ds_name, obj_ds_name = happypose.toolbox.datasets.datasets_cfg.get_obj_ds_info( + cfg.ds_name, + ) + + # drop frames if this was specified + if cfg.n_frames is not None: + scene_ds.frame_index = scene_ds.frame_index[: cfg.n_frames].reset_index( + drop=True, + ) + + # 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) + elif cfg.inference.detection_type == "gt": + detector_model = None + else: + msg = f"Unknown detection_type={cfg.inference.detection_type}" + raise ValueError(msg) + + # Load the coarse and mrefiner models + # Needed to deal with the fact that str and Optional[str] are incompatible types. + # See https://stackoverflow.com/a/53287330 + assert cfg.coarse_run_id is not None + assert cfg.refiner_run_id is not None + # TODO (emaitre): This fuction seems to take the wrong parameters. Trying to fix + # this. + """ + ( + coarse_model, + refiner_model, + mesh_db, + ) = happypose.toolbox.inference.utils.load_pose_models( + coarse_run_id=cfg.coarse_run_id, + refiner_run_id=cfg.refiner_run_id, + n_workers=cfg.n_rendering_workers, + obj_ds_name=obj_ds_name, + urdf_ds_name=urdf_ds_name, + force_panda3d_renderer=True, + ) + """ + object_ds = make_object_dataset(obj_ds_name) + + coarse_model, refiner_model, mesh_db = load_pose_models( + coarse_run_id=cfg.coarse_run_id, + refiner_run_id=cfg.refiner_run_id, + n_workers=8, + ) + + renderer = refiner_model.renderer + + if cfg.inference.run_depth_refiner: + if cfg.inference.depth_refiner == "icp": + ICPRefiner(mesh_db, renderer) + elif cfg.inference.depth_refiner == "teaserpp": + TeaserppRefiner(mesh_db, renderer) + else: + pass + else: + pass + + pose_estimator = PoseEstimator( + refiner_model=refiner_model, + coarse_model=coarse_model, + detector_model=detector_model, + ) + + # Create the prediction runner and run inference + assert cfg.batch_size == 1 + pred_runner = PredictionRunner( + scene_ds=scene_ds, + inference_cfg=cfg.inference, + batch_size=cfg.batch_size, + n_workers=cfg.n_dataloader_workers, + ) + + # Run inference + with torch.no_grad(): + all_preds = pred_runner.get_predictions(pose_estimator) + + logger.info(f"Done with inference on ds={cfg.ds_name}") + logger.info(f"Predictions: {all_preds.keys()}") + + # Keep it simple for now. Only eval the final prediction + eval_keys = set() + eval_keys.add("refiner/final") + eval_keys.add("depth_refiner") + + # Compute eval metrics + # TODO (lmanuelli): Fix this up. + # TODO (ylabbe): Clean this. + eval_metrics, eval_dfs = {}, {} + if not cfg.skip_evaluation: + assert "modelnet" in cfg.ds_name + object_ds = make_object_dataset(obj_ds_name) + mesh_db = MeshDataBase.from_object_ds(object_ds) + meters = { + "modelnet": ModelNetErrorMeter(mesh_db, sample_n_points=None), + } + eval_runner = EvaluationRunner( + scene_ds, + meters, + n_workers=cfg.n_dataloader_workers, + cache_data=False, + batch_size=1, + sampler=pred_runner.sampler, + ) + for preds_k, preds in all_preds.items(): + do_eval = preds_k in set(eval_keys) + if do_eval: + logger.info(f"Evaluation of predictions: {preds_k} (n={len(preds)})") + eval_metrics[preds_k], eval_dfs[preds_k] = eval_runner.evaluate(preds) + else: + logger.info(f"Skipped: {preds_k} (n={len(all_preds)})") + + # Gather predictions from different processes + logger.info("Waiting on barrier.") + torch.distributed.barrier() + logger.info("Gathering predictions from all processes.") + for k, v in all_preds.items(): + all_preds[k] = v.gather_distributed(tmp_dir=get_tmp_dir()).cpu() + + torch.distributed.barrier() + logger.info("Finished gathering predictions from all processes.") + + # Save results to disk + if get_rank() == 0: + results_path = save_dir / "results.pth.tar" + assert cfg.save_dir is not None + save_dir = Path(cfg.save_dir) + save_dir.mkdir(exist_ok=True, parents=True) + logger.info(f"Finished evaluation on {cfg.ds_name}, setting={save_key}") + results = format_results(all_preds, eval_metrics, eval_dfs) + torch.save(results, results_path) + torch.save(results.get("summary"), save_dir / "summary.pth.tar") + torch.save(results.get("predictions"), save_dir / "predictions.pth.tar") + torch.save(results.get("dfs"), save_dir / "error_dfs.pth.tar") + torch.save(results.get("metrics"), save_dir / "metrics.pth.tar") + (save_dir / "summary.txt").write_text(results.get("summary_txt", "")) + (save_dir / "config.yaml").write_text(OmegaConf.to_yaml(cfg)) + logger.info(f"Saved predictions+metrics in {save_dir}") + + return { + "results": results, + "pred_keys": list(all_preds.keys()), + "save_dir": save_dir, + "results_path": results_path, + } + else: + return None diff --git a/happypose/pose_estimators/cosypose/cosypose/training/pose_models_cfg.py b/happypose/pose_estimators/cosypose/cosypose/training/pose_models_cfg.py index 81a71d6f..7fdb28d7 100644 --- a/happypose/pose_estimators/cosypose/cosypose/training/pose_models_cfg.py +++ b/happypose/pose_estimators/cosypose/cosypose/training/pose_models_cfg.py @@ -1,3 +1,8 @@ +from pathlib import Path + +import torch +import yaml + # Backbones from happypose.pose_estimators.cosypose.cosypose.models.efficientnet import EfficientNet from happypose.pose_estimators.cosypose.cosypose.models.flownet import ( @@ -11,6 +16,7 @@ WideResNet34, ) from happypose.pose_estimators.cosypose.cosypose.utils.logging import get_logger +from happypose.toolbox.lib3d.rigid_mesh_database import BatchedMeshes logger = get_logger(__name__) @@ -21,7 +27,7 @@ def check_update_config(config): return config -def create_model_pose(cfg, renderer, mesh_db): +def create_pose_model_cosypose(cfg, renderer, mesh_db): n_inputs = 6 backbone_str = cfg.backbone_str if backbone_str == "efficientnet-b3": @@ -53,9 +59,16 @@ def create_model_pose(cfg, renderer, mesh_db): return model -def create_model_refiner(cfg, renderer, mesh_db): - return create_model_pose(cfg, renderer, mesh_db) - - -def create_model_coarse(cfg, renderer, mesh_db): - return create_model_pose(cfg, renderer, mesh_db) +def load_model_cosypose( + run_dir: Path, renderer, mesh_db_batched: BatchedMeshes, device +): + cfg = yaml.load((run_dir / "config.yaml").read_text(), Loader=yaml.UnsafeLoader) + cfg = check_update_config(cfg) + model = create_pose_model_cosypose(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 diff --git a/happypose/pose_estimators/cosypose/cosypose/training/train_pose.py b/happypose/pose_estimators/cosypose/cosypose/training/train_pose.py index 39163c1a..8de1738a 100644 --- a/happypose/pose_estimators/cosypose/cosypose/training/train_pose.py +++ b/happypose/pose_estimators/cosypose/cosypose/training/train_pose.py @@ -33,6 +33,9 @@ load_pix2pose_results, load_posecnn_results, ) +from happypose.pose_estimators.cosypose.cosypose.training.pose_models_cfg import ( + load_model_cosypose, +) from happypose.pose_estimators.cosypose.cosypose.utils.distributed import ( get_rank, get_world_size, @@ -55,11 +58,13 @@ from happypose.toolbox.renderer.panda3d_batch_renderer import Panda3dBatchRenderer from .pose_forward_loss import h_pose -from .pose_models_cfg import check_update_config, create_model_pose +from .pose_models_cfg import check_update_config, create_pose_model_cosypose cudnn.benchmark = True logger = get_logger(__name__) +device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + def log(config, model, log_dict, test_dict, epoch): save_dir = config.save_dir @@ -93,33 +98,22 @@ def make_eval_bundle(args, model_training): eval_bundle = {} model_training.cfg = args - def load_model(run_id): - if run_id is None: - return None - run_dir = EXP_DIR / run_id - cfg = yaml.load((run_dir / "config.yaml").read_text(), Loader=yaml.FullLoader) - cfg = check_update_config(cfg) - model = ( - create_model_pose( - cfg, - renderer=model_training.renderer, - mesh_db=model_training.mesh_db, - ) - .cuda() - .eval() - ) - ckpt = torch.load(run_dir / "checkpoint.pth.tar")["state_dict"] - model.load_state_dict(ckpt) - model.eval() - model.cfg = cfg - return model - if args.train_refiner: refiner_model = model_training - coarse_model = load_model(args.coarse_run_id_for_test) + coarse_model = load_model_cosypose( + EXP_DIR / args.coarse_run_id_for_test, + model_training.renderer, + model_training.mesh_db, + device, + ) elif args.train_coarse: coarse_model = model_training - refiner_model = load_model(args.refiner_run_id_for_test) + refiner_model = load_model_cosypose( + EXP_DIR / args.refiner_run_id_for_test, + model_training.renderer, + model_training.mesh_db, + device, + ) else: raise ValueError @@ -342,7 +336,9 @@ def make_datasets(dataset_names): .float() ) - model = create_model_pose(cfg=args, renderer=renderer, mesh_db=mesh_db).cuda() + model = create_pose_model_cosypose( + cfg=args, renderer=renderer, mesh_db=mesh_db + ).cuda() eval_bundle = make_eval_bundle(args, model) diff --git a/happypose/pose_estimators/cosypose/cosypose/utils/cosypose_wrapper.py b/happypose/pose_estimators/cosypose/cosypose/utils/cosypose_wrapper.py index 6507ca41..b31f4310 100644 --- a/happypose/pose_estimators/cosypose/cosypose/utils/cosypose_wrapper.py +++ b/happypose/pose_estimators/cosypose/cosypose/utils/cosypose_wrapper.py @@ -8,181 +8,75 @@ """TODO: ---- -- remove commented useless code -- check if all imports necessary -- refactor hardcoded model weight checkpoints. - +- Deprecate this class when possible """ +from typing import Union + import torch -import yaml from happypose.pose_estimators.cosypose.cosypose.config import EXP_DIR -from happypose.pose_estimators.cosypose.cosypose.integrated.detector import Detector 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, -) from happypose.pose_estimators.cosypose.cosypose.training.pose_models_cfg import ( - check_update_config as check_update_config_pose, + load_model_cosypose, ) -from happypose.pose_estimators.cosypose.cosypose.training.pose_models_cfg import ( - create_model_coarse, - create_model_refiner, -) - -# from happypose.pose_estimators.cosypose.cosypose.datasets.datasets_cfg -# import make_scene_dataset, make_object_dataset from happypose.toolbox.datasets.datasets_cfg import make_object_dataset - -# Pose estimator +from happypose.toolbox.datasets.object_dataset import RigidObjectDataset +from happypose.toolbox.inference.utils import load_detector from happypose.toolbox.lib3d.rigid_mesh_database import MeshDataBase from happypose.toolbox.renderer.panda3d_batch_renderer import Panda3dBatchRenderer -""" -def make_object_dataset(example_dir: Path) -> RigidObjectDataset: - print(example_dir) - rigid_objects = [] - mesh_units = "mm" - object_dirs = (example_dir / "meshes").iterdir() - print(object_dirs) - for object_dir in object_dirs: - print(object_dir) - label = object_dir.name - print(label) - mesh_path = None - for fn in object_dir.glob("*"): - print("fn = ", fn) - if fn.suffix in {".obj", ".ply"}: - assert not mesh_path, f"there multiple meshes in the {label} directory" - mesh_path = fn - assert mesh_path, f"couldnt find a obj or ply mesh for {label}" - rigid_objects.append( - RigidObject(label=label, mesh_path=mesh_path, mesh_units=mesh_units) - ) - # TODO: fix mesh units - rigid_object_dataset = RigidObjectDataset(rigid_objects) - return rigid_object_dataset -example_dir = Path("/home/emaitre/cosypose/local_data/bop_datasets/ycbv/examples/") -""" - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") class CosyPoseWrapper: - def __init__(self, dataset_name, n_workers=8, gpu_renderer=False) -> None: + def __init__( + self, + dataset_name: str, + object_dataset: Union[None, RigidObjectDataset] = None, + n_workers=8, + gpu_renderer=False, + ) -> None: self.dataset_name = dataset_name - super().__init__() + self.object_dataset = object_dataset self.detector, self.pose_predictor = self.get_model(dataset_name, n_workers) - @staticmethod - 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 - model = create_model_detector(cfg, len(label_to_category_id)) - 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 - model = Detector(model, ds_name) - return model - - @staticmethod - def load_pose_models(coarse_run_id, refiner_run_id, n_workers): - 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_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_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) - return coarse_model, refiner_model, mesh_db - - @staticmethod - def get_model(dataset_name, n_workers): + def get_model(self, dataset_name, n_workers): # load models - if dataset_name == "tless": + if dataset_name == "hope": + # HOPE setup + # python -m happypose.toolbox.utils.download --cosypose_models=detector-bop-hope-pbr--15246 + # python -m happypose.toolbox.utils.download --cosypose_models=coarse-bop-hope-pbr--225203 + # python -m happypose.toolbox.utils.download --cosypose_models=refiner-bop-hope-pbr--955392 + detector_run_id = "detector-bop-hope-pbr--15246" + coarse_run_id = "coarse-bop-hope-pbr--225203" + refiner_run_id = "refiner-bop-hope-pbr--955392" + elif dataset_name == "tless": # TLESS setup - # python -m cosypose.scripts.download --model=detector-bop-tless-pbr--873074 - # python -m cosypose.scripts.download --model=coarse-bop-tless-pbr--506801 - # python -m cosypose.scripts.download --model=refiner-bop-tless-pbr--233420 + # python -m happypose.toolbox.utils.download --cosypose_models=detector-bop-tless-pbr--873074 + # python -m happypose.toolbox.utils.download --cosypose_models=coarse-bop-tless-pbr--506801 + # python -m happypose.toolbox.utils.download --cosypose_models=refiner-bop-tless-pbr--233420 detector_run_id = "detector-bop-tless-pbr--873074" coarse_run_id = "coarse-bop-tless-pbr--506801" refiner_run_id = "refiner-bop-tless-pbr--233420" elif dataset_name == "ycbv": # YCBV setup - # python -m cosypose.scripts.download --model=detector-bop-ycbv-pbr--970850 - # python -m cosypose.scripts.download --model=coarse-bop-ycbv-pbr--724183 - # python -m cosypose.scripts.download --model=refiner-bop-ycbv-pbr--604090 + # python -m happypose.toolbox.utils.download --cosypose_models=detector-bop-ycbv-pbr--970850 + # python -m happypose.toolbox.utils.download --cosypose_models=coarse-bop-ycbv-pbr--724183 + # python -m happypose.toolbox.utils.download --cosypose_models=refiner-bop-ycbv-pbr--604090 detector_run_id = "detector-bop-ycbv-pbr--970850" coarse_run_id = "coarse-bop-ycbv-pbr--724183" refiner_run_id = "refiner-bop-ycbv-pbr--604090" else: msg = f"Not prepared for {dataset_name} dataset" raise ValueError(msg) - detector = CosyPoseWrapper.load_detector(detector_run_id, dataset_name) - coarse_model, refiner_model, mesh_db = CosyPoseWrapper.load_pose_models( + detector = load_detector(detector_run_id) + coarse_model, refiner_model = self.load_pose_models( coarse_run_id=coarse_run_id, refiner_run_id=refiner_run_id, n_workers=n_workers, @@ -195,6 +89,25 @@ def get_model(dataset_name, n_workers): ) return detector, pose_estimator + def load_pose_models(self, coarse_run_id, refiner_run_id, n_workers): + if self.object_dataset is None: + self.object_dataset = make_object_dataset(self.dataset_name) + renderer = Panda3dBatchRenderer( + self.object_dataset, + n_workers=n_workers, + preload_cache=False, + ) + mesh_db = MeshDataBase.from_object_ds(self.object_dataset) + mesh_db_batched = mesh_db.batched().to(device) + + 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 + def inference(self, observation, coarse_guess=None): detections = None run_detector = True @@ -216,6 +129,5 @@ def inference(self, observation, coarse_guess=None): n_coarse_iterations=0, n_refiner_iterations=4, ) - print("inference successfully.") - # result: this_batch_detections, final_preds + print("inference successfull") return final_preds.cpu() diff --git a/happypose/pose_estimators/cosypose/cosypose/visualization/multiview.py b/happypose/pose_estimators/cosypose/cosypose/visualization/multiview.py index e438eb7a..26d4cdbf 100644 --- a/happypose/pose_estimators/cosypose/cosypose/visualization/multiview.py +++ b/happypose/pose_estimators/cosypose/cosypose/visualization/multiview.py @@ -152,7 +152,7 @@ def make_scene_renderings( {"K": K, "TWC": TWC, "resolution": (w, h)}, ) renders = renderer.render_scene(list_objects, list_cameras) - images = np.stack([render["rgb"] for render in renders]) + images = np.stack([render.rgb for render in renders]) if gui: time.sleep(100) renderer.disconnect() @@ -201,7 +201,7 @@ def render_predictions_wrt_camera(renderer, preds_with_colors, camera): "TWO": preds_with_colors.poses[n].cpu().numpy(), } list_objects.append(obj) - rgb_rendered = renderer.render_scene(list_objects, [camera])[0]["rgb"] + rgb_rendered = renderer.render_scene(list_objects, [camera])[0].rgb return rgb_rendered @@ -212,7 +212,7 @@ def render_gt(renderer, objects, camera, colormap_rgb): obj["color"] = colormap_rgb[obj["label"]] obj["TWO"] = np.linalg.inv(TWC) @ obj["TWO"] camera["TWC"] = np.eye(4) - rgb_rendered = renderer.render_scene(objects, [camera])[0]["rgb"] + rgb_rendered = renderer.render_scene(objects, [camera])[0].rgb return rgb_rendered diff --git a/happypose/pose_estimators/cosypose/cosypose_demos/.gitignore b/happypose/pose_estimators/cosypose/cosypose_demos/.gitignore deleted file mode 100644 index c18dd8d8..00000000 --- a/happypose/pose_estimators/cosypose/cosypose_demos/.gitignore +++ /dev/null @@ -1 +0,0 @@ -__pycache__/ diff --git a/happypose/pose_estimators/cosypose/cosypose_demos/README.md b/happypose/pose_estimators/cosypose/cosypose_demos/README.md deleted file mode 100644 index 82d91bfe..00000000 --- a/happypose/pose_estimators/cosypose/cosypose_demos/README.md +++ /dev/null @@ -1,3 +0,0 @@ -# cosypose_demos -Small wrappers and scripts to run CosyPose -`python3 minimal_test.py` diff --git a/happypose/pose_estimators/megapose/README.md b/happypose/pose_estimators/megapose/README.md index 6c2a6294..fd9c94c4 100644 --- a/happypose/pose_estimators/megapose/README.md +++ b/happypose/pose_estimators/megapose/README.md @@ -275,7 +275,7 @@ This file contains a list of objects with their estimated poses . For each objec Finally, you can visualize the results using: ``` -python -m megapose.scripts.run_inference_on_example barbecue-sauce --vis-outputs +python -m megapose.scripts.run_inference_on_example barbecue-sauce --vis-poses ``` which write several visualization files: diff --git a/happypose/pose_estimators/megapose/inference/detector.py b/happypose/pose_estimators/megapose/inference/detector.py index f621946b..c28339fe 100644 --- a/happypose/pose_estimators/megapose/inference/detector.py +++ b/happypose/pose_estimators/megapose/inference/detector.py @@ -83,6 +83,10 @@ def get_detections( one_instance_per_class: If True, keep only the highest scoring detection within each class. + Returns: + --- + detections: DetectionType=PandasTensorCollection containing bboxes tensor. + bboxes format: xmin, ymin, xmax, ymax """ # [B,3,H,W] @@ -125,28 +129,28 @@ def get_detections( dtype=torch.bool, ).to(device) - outputs = tc.PandasTensorCollection( + detections = tc.PandasTensorCollection( infos=pd.DataFrame(infos), bboxes=bboxes, ) if output_masks: - outputs.register_tensor("masks", masks) + detections.register_tensor("masks", masks) if detection_th is not None: - keep = np.where(outputs.infos["score"] > detection_th)[0] - outputs = outputs[keep] + keep = np.where(detections.infos["score"] > detection_th)[0] + detections = detections[keep] # Keep only the top-detection for each class label if one_instance_per_class: - outputs = happypose.toolbox.inference.utils.filter_detections( - outputs, + detections = happypose.toolbox.inference.utils.filter_detections( + detections, one_instance_per_class=True, ) # Add instance_id column to dataframe # Each detection is now associated with an `instance_id` that # identifies multiple instances of the same object - outputs = happypose.toolbox.inference.utils.add_instance_id(outputs) - return outputs + detections = happypose.toolbox.inference.utils.add_instance_id(detections) + return detections def __call__(self, *args: Any, **kwargs: Any) -> DetectionsType: return self.get_detections(*args, **kwargs) diff --git a/happypose/pose_estimators/megapose/inference/pose_estimator.py b/happypose/pose_estimators/megapose/inference/pose_estimator.py index fe74b4d0..9f46628e 100644 --- a/happypose/pose_estimators/megapose/inference/pose_estimator.py +++ b/happypose/pose_estimators/megapose/inference/pose_estimator.py @@ -18,7 +18,7 @@ # Standard Library import time from collections import defaultdict -from typing import Any, Optional, Tuple +from typing import Any, List, Optional, Tuple # Third Party import numpy as np @@ -37,6 +37,7 @@ ObservationTensor, PoseEstimatesType, ) +from happypose.toolbox.inference.utils import add_instance_id, filter_detections from happypose.toolbox.lib3d.cosypose_ops import TCO_init_from_boxes_autodepth_with_R from happypose.toolbox.utils import transform_utils from happypose.toolbox.utils.logging import get_logger @@ -520,12 +521,12 @@ def run_inference_pipeline( n_refiner_iterations: int = 5, n_pose_hypotheses: int = 1, keep_all_refiner_outputs: bool = False, - detection_filter_kwargs: Optional[dict] = None, run_depth_refiner: bool = False, bsz_images: Optional[int] = None, bsz_objects: Optional[int] = None, cuda_timer: Optional[bool] = False, coarse_estimates: Optional[PoseEstimatesType] = None, + labels_to_keep: List[str] = None, ) -> Tuple[PoseEstimatesType, dict]: """Runs the entire pose estimation pipeline. @@ -567,19 +568,20 @@ def run_inference_pipeline( elapsed = time.time() - start_time timing_str += f"detection={elapsed:.2f}, " - # Ensure that detections has the instance_id column assert detections is not None + # Filter detections + if labels_to_keep is not None: + detections = filter_detections( + detections, + labels_to_keep, + ) + assert ( len(detections) > 0 ), "TOFIX: currently, dealing with absence of detections is not supported" - detections = happypose.toolbox.inference.utils.add_instance_id(detections) - # Filter detections - if detection_filter_kwargs is not None: - detections = happypose.toolbox.inference.utils.filter_detections( - detections, - **detection_filter_kwargs, - ) + # Ensure that detections has the instance_id column + detections = add_instance_id(detections) # Run the coarse estimator using detections data_TCO_coarse, coarse_extra_data = self.forward_coarse_model( diff --git a/happypose/pose_estimators/megapose/scripts/run_inference_on_example.py b/happypose/pose_estimators/megapose/scripts/run_inference_on_example.py index e5eb9c05..2939fc86 100644 --- a/happypose/pose_estimators/megapose/scripts/run_inference_on_example.py +++ b/happypose/pose_estimators/megapose/scripts/run_inference_on_example.py @@ -1,233 +1,63 @@ # Standard Library import argparse -import json import os from pathlib import Path -from typing import List, Tuple, Union +from typing import Dict # Third Party -import numpy as np import torch -from bokeh.io import export_png -from bokeh.plotting import gridplot -from PIL import Image + +from happypose.pose_estimators.megapose.inference.pose_estimator import PoseEstimator # HappyPose -from happypose.toolbox.datasets.object_dataset import RigidObject, RigidObjectDataset -from happypose.toolbox.datasets.scene_dataset import CameraData, ObjectData -from happypose.toolbox.inference.types import ( - DetectionsType, - ObservationTensor, - PoseEstimatesType, +from happypose.toolbox.datasets.object_dataset import RigidObjectDataset +from happypose.toolbox.inference.example_inference_utils import ( + load_detections, + load_object_data, + load_observation_example, + make_detections_visualization, + make_example_object_dataset, + make_poses_visualization, + save_predictions, ) -from happypose.toolbox.inference.utils import make_detections_from_object_data -from happypose.toolbox.lib3d.transform import Transform -from happypose.toolbox.renderer import Panda3dLightData -from happypose.toolbox.renderer.panda3d_scene_renderer import Panda3dSceneRenderer -from happypose.toolbox.utils.conversion import convert_scene_observation_to_panda3d +from happypose.toolbox.inference.types import DetectionsType, ObservationTensor +from happypose.toolbox.inference.utils import filter_detections, load_detector from happypose.toolbox.utils.load_model import NAMED_MODELS, load_named_model from happypose.toolbox.utils.logging import get_logger, set_logging_level -from happypose.toolbox.visualization.bokeh_plotter import BokehPlotter -from happypose.toolbox.visualization.utils import make_contour_overlay - -# MegaPose -# from happypose.toolbox.datasets.object_dataset import RigidObject, RigidObjectDataset -# from happypose.toolbox.datasets.scene_dataset import CameraData, ObjectData - logger = get_logger(__name__) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") -def load_observation( - example_dir: Path, - load_depth: bool = False, -) -> Tuple[np.ndarray, Union[None, np.ndarray], CameraData]: - camera_data = CameraData.from_json((example_dir / "camera_data.json").read_text()) - - rgb = np.array(Image.open(example_dir / "image_rgb.png"), dtype=np.uint8) - assert rgb.shape[:2] == camera_data.resolution - - depth = None - if load_depth: - depth = ( - np.array(Image.open(example_dir / "image_depth.png"), dtype=np.float32) - / 1000 - ) - assert depth.shape[:2] == camera_data.resolution - - return rgb, depth, camera_data - - -def load_observation_tensor( - example_dir: Path, - load_depth: bool = False, -) -> ObservationTensor: - rgb, depth, camera_data = load_observation(example_dir, load_depth) - observation = ObservationTensor.from_numpy(rgb, depth, camera_data.K) - return observation - - -def load_object_data(data_path: Path) -> List[ObjectData]: - object_data = json.loads(data_path.read_text()) - object_data = [ObjectData.from_json(d) for d in object_data] - return object_data - - -def load_detections( - example_dir: Path, -) -> DetectionsType: - input_object_data = load_object_data(example_dir / "inputs/object_data.json") - detections = make_detections_from_object_data(input_object_data).to(device) - return detections - - -def make_object_dataset(example_dir: Path) -> RigidObjectDataset: - rigid_objects = [] - mesh_units = "mm" - object_dirs = (example_dir / "meshes").iterdir() - for object_dir in object_dirs: - label = object_dir.name - mesh_path = None - for fn in object_dir.glob("*"): - if fn.suffix in {".obj", ".ply"}: - assert not mesh_path, f"there multiple meshes in the {label} directory" - mesh_path = fn - assert mesh_path, f"couldnt find a obj or ply mesh for {label}" - rigid_objects.append( - RigidObject(label=label, mesh_path=mesh_path, mesh_units=mesh_units), - ) - # TODO: fix mesh units - rigid_object_dataset = RigidObjectDataset(rigid_objects) - return rigid_object_dataset - - -def make_detections_visualization( - example_dir: Path, -) -> None: - rgb, _, _ = load_observation(example_dir, load_depth=False) - detections = load_detections(example_dir) - plotter = BokehPlotter() - fig_rgb = plotter.plot_image(rgb) - fig_det = plotter.plot_detections(fig_rgb, detections=detections) - output_fn = example_dir / "visualizations" / "detections.png" - output_fn.parent.mkdir(exist_ok=True) - export_png(fig_det, filename=output_fn) - logger.info(f"Wrote detections visualization: {output_fn}") - return - +def setup_pose_estimator(model_name: str, object_dataset: RigidObjectDataset): + logger.info(f"Loading model {model_name}.") + model_info = NAMED_MODELS[model_name] + pose_estimator = load_named_model(model_name, object_dataset).to(device) + # Speed up things by subsampling coarse grid + pose_estimator._SO3_grid = pose_estimator._SO3_grid[::8] -def save_predictions( - example_dir: Path, - pose_estimates: PoseEstimatesType, -) -> None: - labels = pose_estimates.infos["label"] - poses = pose_estimates.poses.cpu().numpy() - object_data = [ - ObjectData(label=label, TWO=Transform(pose)) - for label, pose in zip(labels, poses) - ] - object_data_json = json.dumps([x.to_json() for x in object_data]) - output_fn = example_dir / "outputs" / "object_data.json" - output_fn.parent.mkdir(exist_ok=True) - output_fn.write_text(object_data_json) - logger.info(f"Wrote predictions: {output_fn}") - return + return pose_estimator, model_info def run_inference( - example_dir: Path, - model_name: str, + pose_estimator: PoseEstimator, + model_info: Dict, + observation: ObservationTensor, + detections: DetectionsType, ) -> None: - model_info = NAMED_MODELS[model_name] + observation.to(device) - observation = load_observation_tensor( - example_dir, - load_depth=model_info["requires_depth"], - ) - if torch.cuda.is_available(): - observation.cuda() - detections = load_detections(example_dir).to(device) - object_dataset = make_object_dataset(example_dir) - - logger.info(f"Loading model {model_name}.") - pose_estimator = load_named_model(model_name, object_dataset).to(device) logger.info("Running inference.") - output, _ = pose_estimator.run_inference_pipeline( + data_TCO_final, extra_data = pose_estimator.run_inference_pipeline( observation, detections=detections, **model_info["inference_parameters"], ) + print("Timings:") + print(extra_data["timing_str"]) - save_predictions(example_dir, output) - return - - -def make_output_visualization( - example_dir: Path, -) -> None: - rgb, _, camera_data = load_observation(example_dir, load_depth=False) - camera_data.TWC = Transform(np.eye(4)) - object_datas = load_object_data(example_dir / "outputs" / "object_data.json") - object_dataset = make_object_dataset(example_dir) - - renderer = Panda3dSceneRenderer(object_dataset) - - camera_data, object_datas = convert_scene_observation_to_panda3d( - camera_data, - object_datas, - ) - light_datas = [ - Panda3dLightData( - light_type="ambient", - color=((1.0, 1.0, 1.0, 1)), - ), - ] - renderings = renderer.render_scene( - object_datas, - [camera_data], - light_datas, - render_depth=False, - render_binary_mask=False, - render_normals=False, - copy_arrays=True, - )[0] - - plotter = BokehPlotter() - - fig_rgb = plotter.plot_image(rgb) - fig_mesh_overlay = plotter.plot_overlay(rgb, renderings.rgb) - contour_overlay = make_contour_overlay( - rgb, - renderings.rgb, - dilate_iterations=1, - color=(0, 255, 0), - )["img"] - fig_contour_overlay = plotter.plot_image(contour_overlay) - fig_all = gridplot( - [[fig_rgb, fig_contour_overlay, fig_mesh_overlay]], - toolbar_location=None, - ) - vis_dir = example_dir / "visualizations" - vis_dir.mkdir(exist_ok=True) - export_png(fig_mesh_overlay, filename=vis_dir / "mesh_overlay.png") - export_png(fig_contour_overlay, filename=vis_dir / "contour_overlay.png") - export_png(fig_all, filename=vis_dir / "all_results.png") - logger.info(f"Wrote visualizations to {vis_dir}.") - return - - -# def make_mesh_visualization(RigidObject) -> List[Image]: -# return - - -# def make_scene_visualization(CameraData, List[ObjectData]) -> List[Image]: -# return - - -# def run_inference(example_dir, use_depth: bool = False): -# return + return data_TCO_final.cpu() if __name__ == "__main__": @@ -237,22 +67,52 @@ def make_output_visualization( parser.add_argument( "--model", type=str, - default="megapose-1.0-RGB-multi-hypothesis", + default="megapose-1.0-RGB", ) - parser.add_argument("--vis-detections", action="store_true") + parser.add_argument("--run-detections", action="store_true") parser.add_argument("--run-inference", action="store_true") - parser.add_argument("--vis-outputs", action="store_true") + parser.add_argument("--vis-detections", action="store_true") + parser.add_argument("--vis-poses", action="store_true") args = parser.parse_args() data_dir = os.getenv("HAPPYPOSE_DATA_DIR") - assert data_dir + assert data_dir, "Set HAPPYPOSE_DATA_DIR env variable" example_dir = Path(data_dir) / "examples" / args.example_name + assert ( + example_dir.exists() + ), "Example {args.example_name} not available, follow download instructions" - if args.vis_detections: - make_detections_visualization(example_dir) + # Load data + object_dataset = make_example_object_dataset(example_dir) + rgb, depth, camera_data = load_observation_example(example_dir, load_depth=True) + observation = ObservationTensor.from_numpy(rgb, depth, camera_data.K) + + # Load models + pose_estimator, model_info = setup_pose_estimator(args.model, object_dataset) + + if args.run_detections: + # TODO: hardcoded detector + detector = load_detector(run_id="detector-bop-hope-pbr--15246", device=device) + # Masks are not used for pose prediction, but are computed by Mask-RCNN anyway + detections = detector.get_detections(observation, output_masks=True) + available_labels = [obj.label for obj in object_dataset.list_objects] + detections = filter_detections(detections, available_labels) + else: + detections = load_detections(example_dir).to(device) if args.run_inference: - run_inference(example_dir, args.model) + output = run_inference(pose_estimator, model_info, observation, detections) + save_predictions(output, example_dir) - if args.vis_outputs: - make_output_visualization(example_dir) + if args.vis_detections: + make_detections_visualization(rgb, detections, example_dir) + + if args.vis_poses: + if args.run_inference: + out_filename = "object_data_inf.json" + else: + out_filename = "object_data.json" + object_datas = load_object_data(example_dir / "outputs" / out_filename) + make_poses_visualization( + rgb, object_dataset, object_datas, camera_data, example_dir + ) diff --git a/happypose/toolbox/datasets/datasets_cfg.py b/happypose/toolbox/datasets/datasets_cfg.py index 605bc992..1235e918 100644 --- a/happypose/toolbox/datasets/datasets_cfg.py +++ b/happypose/toolbox/datasets/datasets_cfg.py @@ -249,7 +249,12 @@ def make_scene_dataset( def make_object_dataset(ds_name: str) -> RigidObjectDataset: # BOP original models - if ds_name == "tless.cad": + if ds_name == "tless": + ds: RigidObjectDataset = BOPObjectDataset( + BOP_DS_DIR / "tless/models_cad", + label_format="tless-{label}", + ) + elif ds_name == "tless.cad": ds: RigidObjectDataset = BOPObjectDataset( BOP_DS_DIR / "tless/models_cad", label_format="tless-{label}", diff --git a/happypose/toolbox/inference/example_inference_utils.py b/happypose/toolbox/inference/example_inference_utils.py new file mode 100644 index 00000000..8bf5c6b5 --- /dev/null +++ b/happypose/toolbox/inference/example_inference_utils.py @@ -0,0 +1,174 @@ +import json +from pathlib import Path +from typing import List, Tuple, Union + +import numpy as np +from bokeh.io import export_png +from bokeh.plotting import gridplot +from PIL import Image + +from happypose.toolbox.datasets.object_dataset import RigidObject, RigidObjectDataset +from happypose.toolbox.datasets.scene_dataset import CameraData, ObjectData +from happypose.toolbox.inference.types import DetectionsType, PoseEstimatesType +from happypose.toolbox.inference.utils import make_detections_from_object_data +from happypose.toolbox.lib3d.transform import Transform +from happypose.toolbox.renderer import Panda3dLightData +from happypose.toolbox.renderer.panda3d_scene_renderer import Panda3dSceneRenderer +from happypose.toolbox.utils.conversion import convert_scene_observation_to_panda3d +from happypose.toolbox.utils.logging import get_logger +from happypose.toolbox.visualization.bokeh_plotter import BokehPlotter +from happypose.toolbox.visualization.utils import make_contour_overlay + +logger = get_logger(__name__) + + +def make_example_object_dataset( + example_dir: Path, mesh_units="mm" +) -> RigidObjectDataset: + """ + TODO + """ + + rigid_objects = [] + mesh_units = "mm" + mesh_dir = example_dir / "meshes" + assert mesh_dir.exists(), f"Missing mesh directory {mesh_dir}" + + for mesh_path in mesh_dir.iterdir(): + if mesh_path.suffix in {".obj", ".ply"}: + obj_name = mesh_path.with_suffix("").name + rigid_objects.append( + RigidObject(label=obj_name, mesh_path=mesh_path, mesh_units=mesh_units), + ) + rigid_object_dataset = RigidObjectDataset(rigid_objects) + return rigid_object_dataset + + +def load_observation_example( + example_dir: Path, + load_depth: bool = False, + camera_data_name: str = "camera_data.json", + rgb_name: str = "image_rgb.png", + depth_name: str = "image_depth.png", +) -> Tuple[np.ndarray, Union[None, np.ndarray], CameraData]: + camera_data = CameraData.from_json((example_dir / camera_data_name).read_text()) + + rgb = np.array(Image.open(example_dir / rgb_name), dtype=np.uint8) + assert rgb.shape[:2] == camera_data.resolution + + depth = None + if load_depth: + depth = np.array(Image.open(example_dir / depth_name), dtype=np.float32) / 1000 + assert depth.shape[:2] == camera_data.resolution + + return rgb, depth, camera_data + + +def load_detections( + example_dir: Path, +) -> DetectionsType: + input_object_data = load_object_data(example_dir / "object_data.json") + detections = make_detections_from_object_data(input_object_data) + return detections + + +def load_object_data(data_path: Path) -> List[ObjectData]: + """""" + object_data = json.loads(data_path.read_text()) + object_data = [ObjectData.from_json(d) for d in object_data] + return object_data + + +def make_detections_visualization( + rgb: np.ndarray, + detections: DetectionsType, + example_dir: Path, +) -> None: + plotter = BokehPlotter() + + # TODO: put in BokehPlotter.plot_detections + if hasattr(detections, "masks"): + for mask in detections.masks: + mask = mask.unsqueeze(2).tile((1, 1, 3)).numpy() + rgb[mask] = 122 + + fig_rgb = plotter.plot_image(rgb) + fig_det = plotter.plot_detections(fig_rgb, detections=detections) + output_fn = example_dir / "visualizations" / "detections.png" + output_fn.parent.mkdir(exist_ok=True) + export_png(fig_det, filename=output_fn) + + logger.info(f"Wrote detections visualization: {output_fn}") + return + + +def save_predictions( + pose_estimates: PoseEstimatesType, + example_dir: Path, +) -> None: + labels = pose_estimates.infos["label"] + poses = pose_estimates.poses.cpu().numpy() + object_data = [ + ObjectData(label=label, TWO=Transform(pose)) + for label, pose in zip(labels, poses) + ] + object_data_json = json.dumps([x.to_json() for x in object_data]) + output_fn = example_dir / "outputs" / "object_data_inf.json" + output_fn.parent.mkdir(exist_ok=True) + output_fn.write_text(object_data_json) + logger.info(f"Wrote predictions: {output_fn}") + + +def make_poses_visualization( + rgb: np.ndarray, + object_dataset: RigidObjectDataset, + object_datas: List[ObjectData], + camera_data: CameraData, + example_dir: Path, +) -> None: + camera_data.TWC = Transform(np.eye(4)) + + renderer = Panda3dSceneRenderer(object_dataset) + + camera_data, object_datas = convert_scene_observation_to_panda3d( + camera_data, + object_datas, + ) + light_datas = [ + Panda3dLightData( + light_type="ambient", + color=((1.0, 1.0, 1.0, 1)), + ), + ] + renderings = renderer.render_scene( + object_datas, + [camera_data], + light_datas, + render_depth=False, + render_binary_mask=False, + render_normals=False, + copy_arrays=True, + )[0] + + plotter = BokehPlotter() + + fig_rgb = plotter.plot_image(rgb) + fig_mesh_overlay = plotter.plot_overlay(rgb, renderings.rgb) + contour_overlay = make_contour_overlay( + rgb, + renderings.rgb, + dilate_iterations=1, + color=(0, 255, 0), + )["img"] + fig_contour_overlay = plotter.plot_image(contour_overlay) + fig_all = gridplot( + [[fig_rgb, fig_contour_overlay, fig_mesh_overlay]], + toolbar_location=None, + ) + vis_dir = example_dir / "visualizations" + vis_dir.mkdir(exist_ok=True) + export_png(fig_mesh_overlay, filename=vis_dir / "mesh_overlay.png") + export_png(fig_contour_overlay, filename=vis_dir / "contour_overlay.png") + export_png(fig_all, filename=vis_dir / "all_results.png") + logger.info(f"Wrote visualizations to {vis_dir}.") + return diff --git a/happypose/toolbox/inference/types.py b/happypose/toolbox/inference/types.py index 968ea2a4..3bc71029 100644 --- a/happypose/toolbox/inference/types.py +++ b/happypose/toolbox/inference/types.py @@ -110,9 +110,15 @@ class ObservationTensor: K: Optional[torch.Tensor] = None # [B,3,3] def cuda(self) -> ObservationTensor: - self.images = self.images.cuda() + return self.to("cuda") + + def cpu(self) -> ObservationTensor: + return self.to("cpu") + + def to(self, device): + self.images = self.images.to(device) if self.K is not None: - self.K = self.K.cuda() + self.K = self.K.to(device) return self @property diff --git a/happypose/toolbox/inference/utils.py b/happypose/toolbox/inference/utils.py index f60df533..a99b51c6 100644 --- a/happypose/toolbox/inference/utils.py +++ b/happypose/toolbox/inference/utils.py @@ -25,7 +25,6 @@ from omegaconf import OmegaConf # MegaPose -import happypose.pose_estimators.megapose import happypose.toolbox.utils.tensor_collection as tc from happypose.pose_estimators.megapose.config import EXP_DIR from happypose.pose_estimators.megapose.inference.detector import Detector @@ -48,7 +47,7 @@ from happypose.toolbox.datasets.object_dataset import RigidObjectDataset from happypose.toolbox.datasets.scene_dataset import CameraData, ObjectData from happypose.toolbox.inference.types import DetectionsType, PoseEstimatesType -from happypose.toolbox.lib3d.rigid_mesh_database import MeshDataBase +from happypose.toolbox.lib3d.rigid_mesh_database import BatchedMeshes, MeshDataBase from happypose.toolbox.renderer.panda3d_batch_renderer import Panda3dBatchRenderer from happypose.toolbox.utils.logging import get_logger from happypose.toolbox.utils.models_compat import change_keys_of_older_models @@ -59,13 +58,13 @@ device = torch.device("cuda" if torch.cuda.is_available() else "cpu") -def load_detector(run_id: str) -> torch.nn.Module: +def load_detector(run_id: str, device="cpu") -> torch.nn.Module: run_dir = EXP_DIR / run_id 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 model = create_model_detector(cfg, len(label_to_category_id)) - ckpt = torch.load(run_dir / "checkpoint.pth.tar", map_location=torch.device("cpu")) + ckpt = torch.load(run_dir / "checkpoint.pth.tar", map_location=torch.device(device)) ckpt = ckpt["state_dict"] model.load_state_dict(ckpt) model = model.to(device).eval() @@ -92,7 +91,7 @@ def load_pose_models( ) -> Tuple[ torch.nn.Module, torch.nn.Module, - happypose.toolbox.lib3d.rigid_mesh_database.BatchedMeshes, + BatchedMeshes, ]: coarse_run_dir = models_root / coarse_run_id coarse_cfg: TrainingConfig = load_cfg(coarse_run_dir / "config.yaml") diff --git a/happypose/toolbox/lib3d/multiview.py b/happypose/toolbox/lib3d/multiview.py index 829fc644..ad7e76ee 100644 --- a/happypose/toolbox/lib3d/multiview.py +++ b/happypose/toolbox/lib3d/multiview.py @@ -53,7 +53,7 @@ def _get_views_TCO_pos_sphere(TCO, tCR, cam_positions_wrt_cam0): ref = NodePath("reference_point") ref.reparentTo(root) tWR = TOC[:3, :3] @ tCR.reshape((3, 1)) + TOC[:3, [-1]] - ref.setPos(*tWR[:3]) + ref.setPos(*tWR[:3].flatten()) radius = np.linalg.norm(np.array(tCR)[:3]) cam_positions_wrt_cam0 = cam_positions_wrt_cam0 * radius diff --git a/happypose/toolbox/lib3d/rigid_mesh_database.py b/happypose/toolbox/lib3d/rigid_mesh_database.py index 41d58dec..53e32d19 100644 --- a/happypose/toolbox/lib3d/rigid_mesh_database.py +++ b/happypose/toolbox/lib3d/rigid_mesh_database.py @@ -81,7 +81,7 @@ def from_object_ds(object_ds): obj_list = [object_ds[n] for n in range(len(object_ds))] return MeshDataBase(obj_list) - def batched(self, aabb=False, resample_n_points=None, n_sym=64): + def batched(self, aabb=False, resample_n_points=None, n_sym=64) -> "BatchedMeshes": if aabb: assert resample_n_points is None diff --git a/happypose/toolbox/lib3d/transform.py b/happypose/toolbox/lib3d/transform.py index 8e813e71..9aac7cd1 100644 --- a/happypose/toolbox/lib3d/transform.py +++ b/happypose/toolbox/lib3d/transform.py @@ -132,6 +132,11 @@ def matrix(self) -> np.ndarray: """Returns 4x4 homogeneous matrix representations.""" return self._T.homogeneous + @property + def tensor(self) -> np.ndarray: + """Returns 4x4 homogeneous matrix representations.""" + return torch.tensor(self._T.homogeneous) + @staticmethod def Identity(): return Transform(pin.SE3.Identity()) diff --git a/happypose/toolbox/renderer/bullet_scene_renderer.py b/happypose/toolbox/renderer/bullet_scene_renderer.py index 40d3213f..5d8cbc44 100644 --- a/happypose/toolbox/renderer/bullet_scene_renderer.py +++ b/happypose/toolbox/renderer/bullet_scene_renderer.py @@ -7,6 +7,7 @@ # from happypose.toolbox.datasets.datasets_cfg import UrdfDataset from happypose.pose_estimators.cosypose.cosypose.datasets.urdf_dataset import ( + OneUrdfDataset, UrdfDataset, ) @@ -31,7 +32,7 @@ def __init__( gpu_renderer=True, gui=False, ): - if isinstance(asset_dataset, UrdfDataset): + if isinstance(asset_dataset, (UrdfDataset, OneUrdfDataset)): self.urdf_ds = asset_dataset elif isinstance(asset_dataset, RigidObjectDataset): # Build urdfs files from RigidObjectDataset @@ -43,7 +44,7 @@ def __init__( 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" + f"asset_dataset of type {type(asset_dataset)} should be either OneUrdfDataset/UrdfDataset or RigidObjectDataset" ) self.connect(gpu_renderer=gpu_renderer, gui=gui) self.body_cache = BodyCache(self.urdf_ds, self.client_id) diff --git a/happypose/toolbox/renderer/types.py b/happypose/toolbox/renderer/types.py index abae1b4a..6d960c64 100644 --- a/happypose/toolbox/renderer/types.py +++ b/happypose/toolbox/renderer/types.py @@ -160,6 +160,10 @@ class Panda3dObjectData: scale: float = 1 positioning_function: Optional[NodeFunction] = None + def __post_init__(self): + if not isinstance(self.TWO, Transform): + self.TWO = Transform(self.TWO) + def set_node_material_and_transparency( self, node_path: p3d.core.NodePath, diff --git a/happypose/toolbox/utils/download.py b/happypose/toolbox/utils/download.py index 8ca2e210..d95e60d5 100755 --- a/happypose/toolbox/utils/download.py +++ b/happypose/toolbox/utils/download.py @@ -55,13 +55,16 @@ "tudl": { "splits": ["test_all", "train_real"], }, + "hope": { + "splits": ["test_all", "train_real"], + }, } BOP_DS_NAMES = list(BOP_DATASETS.keys()) async def main(): - parser = argparse.ArgumentParser("CosyPose download utility") + parser = argparse.ArgumentParser("HappyPose download utility") parser.add_argument("--bop_dataset", nargs="*", choices=BOP_DS_NAMES) parser.add_argument("--bop_extra_files", nargs="*", choices=["ycbv", "tless"]) parser.add_argument("--cosypose_models", nargs="*") diff --git a/tests/data/bullet_obj_000001_batch_render.png b/tests/data/bullet_obj_000001_batch_render.png new file mode 100644 index 00000000..517a7530 Binary files /dev/null and b/tests/data/bullet_obj_000001_batch_render.png differ diff --git a/tests/data/bullet_obj_000001_scene_render.png b/tests/data/bullet_obj_000001_scene_render.png new file mode 100644 index 00000000..f8cc36b8 Binary files /dev/null and b/tests/data/bullet_obj_000001_scene_render.png differ diff --git a/tests/data/panda3d_obj_batch_render.png b/tests/data/panda3d_obj_batch_render.png new file mode 100644 index 00000000..885b64fb Binary files /dev/null and b/tests/data/panda3d_obj_batch_render.png differ diff --git a/tests/data/panda3d_obj_scene_render.png b/tests/data/panda3d_obj_scene_render.png new file mode 100644 index 00000000..401cb8af Binary files /dev/null and b/tests/data/panda3d_obj_scene_render.png differ diff --git a/tests/test_cosypose_inference.py b/tests/test_cosypose_inference.py index 238bdb1c..6eb9a45d 100644 --- a/tests/test_cosypose_inference.py +++ b/tests/test_cosypose_inference.py @@ -4,31 +4,21 @@ import numpy as np import pinocchio as pin -import torch -import yaml -from PIL import Image from happypose.pose_estimators.cosypose.cosypose.config import EXP_DIR, LOCAL_DATA_DIR -from happypose.pose_estimators.cosypose.cosypose.integrated.detector import Detector from happypose.pose_estimators.cosypose.cosypose.integrated.pose_estimator import ( PoseEstimator, ) -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.training.pose_models_cfg import ( - check_update_config as check_update_config_pose, + load_model_cosypose, ) -from happypose.pose_estimators.cosypose.cosypose.training.pose_models_cfg import ( - create_model_coarse, - create_model_refiner, +from happypose.toolbox.datasets.bop_object_datasets import ( + RigidObject, + RigidObjectDataset, ) -from happypose.toolbox.datasets.bop_object_datasets import BOPObjectDataset -from happypose.toolbox.datasets.scene_dataset import CameraData +from happypose.toolbox.inference.example_inference_utils import load_observation_example from happypose.toolbox.inference.types import ObservationTensor +from happypose.toolbox.inference.utils import load_detector from happypose.toolbox.lib3d.rigid_mesh_database import MeshDataBase from happypose.toolbox.renderer.panda3d_batch_renderer import Panda3dBatchRenderer @@ -36,133 +26,74 @@ class TestCosyPoseInference(unittest.TestCase): """Unit tests for CosyPose inference example.""" - @staticmethod - def _load_detector( - device="cpu", - ds_name="ycbv", - run_id="detector-bop-ycbv-pbr--970850", - ): - """Load CosyPose detector.""" - run_dir = EXP_DIR / run_id - assert run_dir.exists(), "The run_id is invalid, or you forget to download data" - cfg = check_update_config_detector( - yaml.load((run_dir / "config.yaml").read_text(), Loader=yaml.UnsafeLoader), - ) - label_to_category_id = cfg.label_to_category_id - ckpt = torch.load(run_dir / "checkpoint.pth.tar", map_location=device)[ - "state_dict" - ] - model = create_model_detector(cfg, len(label_to_category_id)) - model.load_state_dict(ckpt) - model = model.to(device).eval() - model.cfg = cfg - model.config = cfg - return Detector(model, ds_name) - - @staticmethod - def _load_pose_model(run_id, renderer, mesh_db, device): - """Load either coarse or refiner model (decided based on run_id/config).""" - run_dir = EXP_DIR / run_id - cfg = yaml.load((run_dir / "config.yaml").read_text(), Loader=yaml.UnsafeLoader) - cfg = check_update_config_pose(cfg) - - f_mdl = create_model_refiner if cfg.train_refiner else create_model_coarse - ckpt = torch.load(run_dir / "checkpoint.pth.tar", map_location=device)[ - "state_dict" - ] - model = f_mdl(cfg, renderer=renderer, mesh_db=mesh_db) - model.load_state_dict(ckpt) - model = model.to(device).eval() - model.cfg = cfg - model.config = cfg - return model - - @staticmethod - def _load_pose_models( - coarse_run_id="coarse-bop-ycbv-pbr--724183", - refiner_run_id="refiner-bop-ycbv-pbr--604090", - n_workers=1, - device="cpu", - ): - """Load coarse and refiner for the crackers example renderer.""" - object_dataset = BOPObjectDataset( - LOCAL_DATA_DIR / "examples" / "crackers_example" / "models", - label_format="ycbv-{label}", + def test_cosypose_pipeline(self): + """Run detector with coarse and refiner from CosyPose.""" + expected_object_label = "hope-obj_000002" + mesh_file_name = "hope-obj_000002.ply" + data_dir = LOCAL_DATA_DIR / "examples" / "barbecue-sauce" + mesh_dir = data_dir / "meshes" + mesh_path = mesh_dir / mesh_file_name + device = "cpu" + n_workers = 1 + + rgb, depth, camera_data = load_observation_example(data_dir, load_depth=True) + # TODO: cosypose forward does not work if depth is loaded detection contrary to megapose + observation = ObservationTensor.from_numpy(rgb, depth=None, K=camera_data.K) + + detector = load_detector(run_id="detector-bop-hope-pbr--15246", device=device) + # detections = detector.get_detections(observation=observation) + + object_dataset = RigidObjectDataset( + objects=[ + RigidObject( + label=expected_object_label, mesh_path=mesh_path, mesh_units="mm" + ) + ] ) + renderer = Panda3dBatchRenderer( object_dataset, n_workers=n_workers, preload_cache=False, ) + coarse_run_id = "coarse-bop-hope-pbr--225203" + refiner_run_id = "refiner-bop-hope-pbr--955392" + mesh_db = MeshDataBase.from_object_ds(object_dataset) - mesh_db_batched = mesh_db.batched().to(device) - kwargs = {"renderer": renderer, "mesh_db": mesh_db_batched, "device": device} - coarse_model = TestCosyPoseInference._load_pose_model(coarse_run_id, **kwargs) - refiner_model = TestCosyPoseInference._load_pose_model(refiner_run_id, **kwargs) - return coarse_model, refiner_model - - @staticmethod - def _load_crackers_example_observation(): - """Load cracker example observation tensor.""" - data_dir = LOCAL_DATA_DIR / "examples" / "crackers_example" - camera_data = CameraData.from_json((data_dir / "camera_data.json").read_text()) - rgb = np.array(Image.open(data_dir / "image_rgb.png"), dtype=np.uint8) - assert rgb.shape[:2] == camera_data.resolution - return ObservationTensor.from_numpy(rgb=rgb, K=camera_data.K) - - def test_detector(self): - """Run detector on known image to see if cracker box is detected.""" - observation = self._load_crackers_example_observation() - detector = self._load_detector() - detections = detector.get_detections(observation=observation) - for s1, s2 in zip(detections.infos.score, detections.infos.score[1:]): - self.assertGreater(s1, s2) # checks that observations are ordered - - self.assertGreater(len(detections), 0) - self.assertEqual(detections.infos.label[0], "ycbv-obj_000002") - self.assertGreater(detections.infos.score[0], 0.8) - - xmin, ymin, xmax, ymax = detections.bboxes[0] - # assert expected obj center inside BB - self.assertTrue(xmin < 320 < xmax and ymin < 250 < ymax) - # assert a few outside points are outside BB - self.assertFalse(xmin < 100 < xmax and ymin < 50 < ymax) - self.assertFalse(xmin < 300 < xmax and ymin < 50 < ymax) - self.assertFalse(xmin < 500 < xmax and ymin < 50 < ymax) - self.assertFalse(xmin < 100 < xmax and ymin < 250 < ymax) - self.assertTrue(xmin < 300 < xmax and ymin < 250 < ymax) - self.assertFalse(xmin < 500 < xmax and ymin < 250 < ymax) - self.assertFalse(xmin < 100 < xmax and ymin < 450 < ymax) - self.assertFalse(xmin < 300 < xmax and ymin < 450 < ymax) - self.assertFalse(xmin < 500 < xmax and ymin < 450 < ymax) + mesh_db_batched = mesh_db.batched().to("cpu") + 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 + ) - def test_cosypose_pipeline(self): - """Run detector with coarse and refiner.""" - observation = self._load_crackers_example_observation() - detector = self._load_detector() - coarse_model, refiner_model = self._load_pose_models() pose_estimator = PoseEstimator( refiner_model=refiner_model, coarse_model=coarse_model, detector_model=detector, ) + + # Run detector and pose estimator filtering object labels preds, _ = pose_estimator.run_inference_pipeline( observation=observation, detection_th=0.8, run_detector=True, + n_refiner_iterations=3, + labels_to_keep=[expected_object_label], ) self.assertEqual(len(preds), 1) - self.assertEqual(preds.infos.label[0], "ycbv-obj_000002") + self.assertEqual(preds.infos.label[0], expected_object_label) pose = pin.SE3(preds.poses[0].numpy()) exp_pose = pin.SE3( - pin.exp3(np.array([1.44, 1.19, -0.91])), - np.array([0, 0, 0.52]), + pin.exp3(np.array([1.4, 1.6, -1.11])), + np.array([0.1, 0.07, 0.45]), ) diff = pose.inverse() * exp_pose - self.assertLess(np.linalg.norm(pin.log6(diff).vector), 0.1) + self.assertLess(np.linalg.norm(pin.log6(diff).vector), 0.3) if __name__ == "__main__": diff --git a/tests/test_detector.py b/tests/test_detector.py new file mode 100644 index 00000000..46f0fbf7 --- /dev/null +++ b/tests/test_detector.py @@ -0,0 +1,56 @@ +"""Set of unit tests for testing inference example for CosyPose.""" +import unittest + +from happypose.pose_estimators.cosypose.cosypose.config import LOCAL_DATA_DIR +from happypose.toolbox.inference.example_inference_utils import load_observation_example +from happypose.toolbox.inference.types import ObservationTensor +from happypose.toolbox.inference.utils import filter_detections, load_detector + + +class TestDetector(unittest.TestCase): + """Unit tests for CosyPose inference example.""" + + def test_detector(self): + """Run detector on known image to see if object is detected.""" + + expected_object_label = "hope-obj_000002" + data_dir = LOCAL_DATA_DIR / "examples" / "barbecue-sauce" + + rgb, depth, camera_data = load_observation_example(data_dir, load_depth=True) + # TODO: cosypose forward does not work if depth is loaded detection contrary to megapose + observation = ObservationTensor.from_numpy(rgb, depth=None, K=camera_data.K) + + detector = load_detector(run_id="detector-bop-hope-pbr--15246", device="cpu") + detections = detector.get_detections(observation=observation) + detections = filter_detections(detections, labels=[expected_object_label]) + + for s1, s2 in zip(detections.infos.score, detections.infos.score[1:]): + self.assertGreater(s1, s2) # checks that observations are ordered + + self.assertGreater(len(detections), 0) + self.assertEqual(detections.infos.label[0], expected_object_label) + self.assertGreater(detections.infos.score[0], 0.8) + + bbox = detections.bboxes[0] # xmin, ymin, xmax, ymax + + # assert if different sample points lie inside/outside the bounding box + self.assertTrue(is_point_in([430, 300], bbox)) + self.assertTrue(is_point_in([430, 400], bbox)) + self.assertTrue(is_point_in([490, 255], bbox)) + + self.assertFalse(is_point_in([0, 0], bbox)) + self.assertFalse(is_point_in([200, 400], bbox)) + self.assertFalse(is_point_in([490, 50], bbox)) + self.assertFalse(is_point_in([550, 500], bbox)) + + +def is_point_in(p, bb): + """ + p: pixel sequence (x,y) + bb: bounding box sequence (xmin, ymin, xmax, ymax) + """ + return bb[0] < p[0] < bb[2] and bb[1] < p[1] < bb[3] + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_lib3d.py b/tests/test_lib3d.py index edbc3c5b..16e23088 100644 --- a/tests/test_lib3d.py +++ b/tests/test_lib3d.py @@ -4,6 +4,7 @@ import pinocchio as pin import torch +# from numpy.testing import assert_equal as np.allclose from happypose.toolbox.lib3d.transform import Transform @@ -19,12 +20,11 @@ def test_constructor(self): # 1 arg constructor T1 = Transform(M1) - assert T1._T == M1 - + self.assertTrue(T1._T == M1) T1b = Transform(T1) T1c = Transform(M1.homogeneous) T1d = Transform(torch.from_numpy(M1.homogeneous)) - assert T1 == T1b == T1c == T1d + self.assertTrue(T1 == T1b == T1c == T1d) # 2 args constructor R1 = M1.rotation @@ -38,21 +38,24 @@ def test_constructor(self): T1i = Transform(list(pin.Quaternion(R1).coeffs()), t1) T1j = Transform(torch.from_numpy(pin.Quaternion(R1).coeffs().copy()), t1) - assert T1 == T1e == T1f - assert np.all( - np.isclose(T1._T, T1g._T) - ) # small error due to Quaternion conversion back and forth - assert T1g == T1h == T1h == T1i == T1j + self.assertTrue(T1 == T1e == T1f) + self.assertTrue(np.allclose(T1._T, T1g._T)) + self.assertTrue(T1g == T1h == T1h == T1i == T1j) + + # Conversions + self.assertTrue(np.allclose(M1.homogeneous, T1.toHomogeneousMatrix())) + self.assertTrue(np.allclose(M1.homogeneous, T1.matrix)) + self.assertTrue(torch.allclose(torch.from_numpy(M1.homogeneous), T1.tensor)) + # Composition T2 = Transform(M2) T3 = Transform(M3) T3m = T1 * T2 - assert T3 == T3m + self.assertTrue(T3 == T3m) + # Inverse T1inv = Transform(T1.inverse()) - assert T1inv == T1.inverse() - - assert np.all(np.isclose(T1.toHomogeneousMatrix(), M1.homogeneous)) + self.assertTrue(T1inv == T1.inverse()) if __name__ == "__main__": diff --git a/tests/test_megapose_inference.py b/tests/test_megapose_inference.py index b53a39a5..d15b4ccb 100644 --- a/tests/test_megapose_inference.py +++ b/tests/test_megapose_inference.py @@ -6,50 +6,66 @@ import pinocchio as pin from happypose.pose_estimators.cosypose.cosypose.config import LOCAL_DATA_DIR -from happypose.toolbox.datasets.bop_object_datasets import BOPObjectDataset +from happypose.toolbox.datasets.bop_object_datasets import ( + RigidObject, + RigidObjectDataset, +) +from happypose.toolbox.inference.example_inference_utils import load_observation_example +from happypose.toolbox.inference.types import ObservationTensor +from happypose.toolbox.inference.utils import load_detector from happypose.toolbox.utils.load_model import NAMED_MODELS, load_named_model -from .test_cosypose_inference import TestCosyPoseInference - class TestMegaPoseInference(unittest.TestCase): """Unit tests for MegaPose inference example.""" def test_megapose_pipeline(self): - """Run detector from CosyPose with coarse and refiner from MegaPose.""" - observation = TestCosyPoseInference._load_crackers_example_observation() - - detector = TestCosyPoseInference._load_detector() - detections = detector.get_detections(observation=observation) - - self.assertGreater(len(detections), 0) - detections = detections[:1] # let's keep the most promising one only. - - object_dataset = BOPObjectDataset( - LOCAL_DATA_DIR / "examples" / "crackers_example" / "models", - label_format="ycbv-{label}", + """Run detector from with coarse and refiner from MegaPose.""" + expected_object_label = "hope-obj_000002" + mesh_file_name = "hope-obj_000002.ply" + data_dir = LOCAL_DATA_DIR / "examples" / "barbecue-sauce" + mesh_dir = data_dir / "meshes" + mesh_path = mesh_dir / mesh_file_name + + rgb, depth, camera_data = load_observation_example(data_dir, load_depth=True) + # TODO: cosypose forward does not work if depth is loaded detection contrary to megapose + observation = ObservationTensor.from_numpy(rgb, depth=None, K=camera_data.K) + + detector = load_detector(run_id="detector-bop-hope-pbr--15246", device="cpu") + + object_dataset = RigidObjectDataset( + objects=[ + RigidObject( + label=expected_object_label, mesh_path=mesh_path, mesh_units="mm" + ) + ] ) - model_info = NAMED_MODELS["megapose-1.0-RGB"] - pose_estimator = load_named_model("megapose-1.0-RGB", object_dataset).to("cpu") - # let's limit the grid, 278 is the most promising one, 477 the least one - pose_estimator._SO3_grid = pose_estimator._SO3_grid[[278, 477]] + model_name = "megapose-1.0-RGB" + model_info = NAMED_MODELS[model_name] + pose_estimator = load_named_model(model_name, object_dataset).to("cpu") + # Uniformely sumbsample the grid to increase speed + pose_estimator._SO3_grid = pose_estimator._SO3_grid[::8] + pose_estimator.detector_model = detector + + # Run detector and pose estimator filtering object labels preds, data = pose_estimator.run_inference_pipeline( observation, - detections=detections, + run_detector=True, **model_info["inference_parameters"], + labels_to_keep=[expected_object_label], ) scores = data["coarse"]["data"]["logits"] - self.assertGreater(scores[0, 0], scores[0, 1]) # 278 is better than 477 + self.assertGreater(scores[0, 0], scores[0, 1]) self.assertEqual(len(preds), 1) - self.assertEqual(preds.infos.label[0], "ycbv-obj_000002") + self.assertEqual(preds.infos.label[0], expected_object_label) pose = pin.SE3(preds.poses[0].numpy()) exp_pose = pin.SE3( - pin.exp3(np.array([1.44, 1.19, -0.91])), - np.array([0, 0, 0.52]), + pin.exp3(np.array([1.4, 1.6, -1.11])), + np.array([0.1, 0.07, 0.45]), ) diff = pose.inverse() * exp_pose self.assertLess(np.linalg.norm(pin.log6(diff).vector), 0.3)