From de16a7826d5d795b33f9f6580e349b6adfae057c Mon Sep 17 00:00:00 2001
From: Vladimir Petrik <vladimir.petrik@cvut.cz>
Date: Mon, 18 Sep 2023 10:53:38 +0200
Subject: [PATCH] test megapose inference added

---
 .github/workflows/test.yml                 |  2 +
 happypose/toolbox/utils/transform_utils.py |  3 +-
 tests/test_cosypose_inference.py           |  5 ++-
 tests/test_megapose_inference.py           | 49 ++++++++++++++++++++++
 4 files changed, 56 insertions(+), 3 deletions(-)
 create mode 100644 tests/test_megapose_inference.py

diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
index ae5224c7..285790fa 100644
--- a/.github/workflows/test.yml
+++ b/.github/workflows/test.yml
@@ -71,6 +71,8 @@ jobs:
           python -m happypose.toolbox.utils.download --cosypose_model=coarse-bop-ycbv-pbr--724183
           python -m happypose.toolbox.utils.download --cosypose_model=refiner-bop-ycbv-pbr--604090
           
+          python -m happypose.toolbox.utils.download --megapose_models
+          
           cd tests/data
           git clone https://github.com/petrikvladimir/happypose_test_data.git crackers_example
 
diff --git a/happypose/toolbox/utils/transform_utils.py b/happypose/toolbox/utils/transform_utils.py
index c1741716..411a46d8 100644
--- a/happypose/toolbox/utils/transform_utils.py
+++ b/happypose/toolbox/utils/transform_utils.py
@@ -34,7 +34,8 @@ def load_SO3_grid(resolution):
     Returns:
         rotmats: [N,3,3]
     """
-    data_fname = PROJECT_DIR / f"src/megapose/data/data_{resolution}.qua"
+    meagpose_dir = PROJECT_DIR / "happypose" / "pose_estimators" / "megapose"
+    data_fname = meagpose_dir / f"src/megapose/data/data_{resolution}.qua"
 
     assert data_fname.is_file(), f"File {data_fname} not found"
 
diff --git a/tests/test_cosypose_inference.py b/tests/test_cosypose_inference.py
index c204edcc..d32f929c 100644
--- a/tests/test_cosypose_inference.py
+++ b/tests/test_cosypose_inference.py
@@ -97,12 +97,13 @@ def _load_pose_models(
         refiner_model = TestCosyPoseInference._load_pose_model(refiner_run_id, **kwargs)
         return coarse_model, refiner_model
 
-    def _load_crackers_example_observation(self):
+    @staticmethod
+    def _load_crackers_example_observation():
         """Load cracker example observation tensor"""
         data_dir = Path(__file__).parent.joinpath("data").joinpath("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)
-        self.assertEqual(rgb.shape[:2], camera_data.resolution)
+        assert rgb.shape[:2] == camera_data.resolution
         return ObservationTensor.from_numpy(rgb=rgb, K=camera_data.K)
 
     def test_detector(self):
diff --git a/tests/test_megapose_inference.py b/tests/test_megapose_inference.py
new file mode 100644
index 00000000..1a1b0ecd
--- /dev/null
+++ b/tests/test_megapose_inference.py
@@ -0,0 +1,49 @@
+"""Set of unit tests for testing inference example for MegaPose."""
+import unittest
+
+import numpy as np
+import pinocchio as pin
+from pathlib import Path
+
+from .test_cosypose_inference import TestCosyPoseInference
+from happypose.toolbox.datasets.bop_object_datasets import BOPObjectDataset
+from happypose.toolbox.utils.load_model import NAMED_MODELS, load_named_model
+
+
+class TestMegaPoseInference(unittest.TestCase):
+    """Unit tests for MegaPose inference example."""
+
+    def test_meagpose_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(
+            Path(__file__).parent / "data" / "crackers_example" / "models",
+            label_format="ycbv-{label}",
+        )
+
+        model_info = NAMED_MODELS["megapose-1.0-RGB"]
+        pose_estimator = load_named_model("megapose-1.0-RGB", object_dataset).to("cpu")
+        preds, _ = pose_estimator.run_inference_pipeline(
+            observation, detections=detections, **model_info["inference_parameters"]
+        )
+
+        self.assertEqual(len(preds), 1)
+        self.assertEqual(preds.infos.label[0], "ycbv-obj_000002")
+
+        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])
+        )
+        diff = pose.inverse() * exp_pose
+        self.assertLess(np.linalg.norm(pin.log6(diff).vector), 0.2)
+
+
+if __name__ == "__main__":
+    unittest.main()