Merge branch '90-coarse-evaluation' into next-pack

happypose/pose_estimators/megapose/evaluation/bop.py

@@ -38,6 +38,10 @@
 from happypose.pose_estimators.megapose.evaluation.eval_config import BOPEvalConfig
 from happypose.toolbox.datasets.scene_dataset import ObjectData
 from happypose.toolbox.inference.utils import make_detections_from_object_data
+from happypose.toolbox.utils.tensor_collection import (
+    PandasTensorCollection,
+    filter_top_pose_estimates,
+)
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
@@ -146,14 +150,41 @@ def convert_results_to_coco(results_path, out_json_path, detection_method):
     return
 
 
+def get_best_coarse_predictions(coarse_preds: PandasTensorCollection):
+    group_cols = ["scene_id", "view_id", "label", "instance_id"]
+    coarse_preds = filter_top_pose_estimates(
+        coarse_preds,
+        top_K=1,
+        group_cols=group_cols,
+        filter_field="coarse_score",
+        ascending=False,
+    )
+    coarse_preds.infos = coarse_preds.infos.rename(
+        columns={"coarse_score": "pose_score"}
+    )
+    return coarse_preds
+
+
 def convert_results_to_bop(
     results_path: Path,
     out_csv_path: Path,
     method: str,
     use_pose_score: bool = True,
 ):
+    """
+    results_path: path to file storing a pickled dictionary,
+                  with a "predictions" key storing all results of a given evaluation
+    out_csv_path: path where bop format csv is saved
+    method: key to one of the available method predictions
+    use_pose_score: if true, uses the score obtained from the pose estimator, otherwise
+                    from the detector.
+    """
+
     predictions = torch.load(results_path)["predictions"]
     predictions = predictions[method]
+    if method == "coarse":
+        predictions = get_best_coarse_predictions(predictions)
+
     print("Predictions from:", results_path)
     print("Method:", method)
     print("Number of predictions: ", len(predictions))

happypose/pose_estimators/megapose/evaluation/eval_config.py

@@ -88,7 +88,8 @@ class FullEvalConfig(EvalConfig):
     detection_coarse_types: Optional[list] = None
     ds_names: Optional[list[str]] = None
     run_bop_eval: bool = True
-    modelnet_categories: Optional[list[str]] = None
+    eval_coarse_also: bool = False
+    convert_only: bool = False
 
 
 @dataclass

happypose/pose_estimators/megapose/evaluation/prediction_runner.py

@@ -15,7 +15,6 @@
 
 
 # Standard Library
-import time
 from collections import defaultdict
 from typing import Optional
 
@@ -128,7 +127,6 @@ def run_inference_pipeline(
             coarse_estimates.infos["instance_id"] = 0
             run_detector = False
 
-        t = time.time()
         preds, extra_data = pose_estimator.run_inference_pipeline(
             obs_tensor,
             detections=detections,
@@ -140,7 +138,6 @@ def run_inference_pipeline(
             bsz_images=self.inference_cfg.bsz_images,
             bsz_objects=self.inference_cfg.bsz_objects,
         )
-        time.time() - t
 
         # TODO (lmanuelli): Process this into a dict with keys like
         # - 'refiner/iteration=1`
@@ -150,26 +147,33 @@ def run_inference_pipeline(
         # go back and extract out the 'refiner/iteration=1`, `refiner/iteration=5`
         # things for the ones that were actually the highest scoring at the end.
 
-        all_preds = {}
-        data_TCO_refiner = extra_data["refiner"]["preds"]
-
         ref_str = f"refiner/iteration={self.inference_cfg.n_refiner_iterations}"
         all_preds = {
             "final": preds,
-            ref_str: data_TCO_refiner,
-            "refiner/final": data_TCO_refiner,
+            ref_str: extra_data["refiner"]["preds"],
+            "refiner/final": extra_data["refiner"]["preds"],
             "coarse": extra_data["coarse"]["preds"],
+            "coarse_filter": extra_data["coarse_filter"]["preds"],
+        }
+
+        # Only keep necessary metadata
+        del extra_data["coarse"]["data"]["TCO"]
+        all_preds_data = {
+            "coarse": extra_data["coarse"]["data"],
+            "refiner": extra_data["refiner"]["data"],
+            "scoring": extra_data["scoring"],
         }
 
         if self.inference_cfg.run_depth_refiner:
             all_preds["depth_refiner"] = extra_data["depth_refiner"]["preds"]
+            all_preds_data["depth_refiner"] = extra_data["depth_refiner"]["data"]
 
         for _k, v in all_preds.items():
             if "mask" in v.tensors:
                 # breakpoint()
                 v.delete_tensor("mask")
 
-        return all_preds
+        return all_preds, all_preds_data
 
     def get_predictions(
         self,
@@ -246,7 +250,7 @@ def get_predictions(
             cuda_timer = CudaTimer()
             cuda_timer.start()
             with torch.no_grad():
-                all_preds = self.run_inference_pipeline(
+                all_preds, all_preds_data = self.run_inference_pipeline(
                     pose_estimator,
                     obs_tensor,
                     gt_detections,
@@ -256,18 +260,50 @@ def get_predictions(
             cuda_timer.end()
             duration = cuda_timer.elapsed()
 
-            total_duration = duration + dt_det
+            duration + dt_det
 
             # Add metadata to the predictions for later evaluation
-            for k, v in all_preds.items():
-                v.infos["time"] = total_duration
-                v.infos["scene_id"] = scene_id
-                v.infos["view_id"] = view_id
-                predictions_list[k].append(v)
+            for pred_name, pred in all_preds.items():
+                pred.infos["time"] = dt_det + compute_pose_est_total_time(
+                    all_preds_data,
+                    pred_name,
+                )
+                pred.infos["scene_id"] = scene_id
+                pred.infos["view_id"] = view_id
+                predictions_list[pred_name].append(pred)
 
         # Concatenate the lists of PandasTensorCollections
         predictions = {}
         for k, v in predictions_list.items():
             predictions[k] = tc.concatenate(v)
 
         return predictions
+
+
+def compute_pose_est_total_time(all_preds_data: dict, pred_name: str):
+    # all_preds_data:
+    # dict_keys(
+    # ["final", "refiner/iteration=5", "refiner/final", "coarse", "coarse_filter"]
+    # )  # optionally 'depth_refiner'
+    dt_coarse = all_preds_data["coarse"]["time"]
+    dt_coarse_refiner = dt_coarse + all_preds_data["refiner"]["time"]
+    if "depth_refiner" in all_preds_data:
+        dt_coarse_refiner_depth = (
+            dt_coarse_refiner + all_preds_data["depth_refiner"]["time"]
+        )
+
+    if pred_name.startswith("coarse"):
+        return dt_coarse
+    elif pred_name.startswith("refiner"):
+        return dt_coarse_refiner
+    elif pred_name == "depth_refiner":
+        return dt_coarse_refiner_depth
+    elif pred_name == "final":
+        return (
+            dt_coarse_refiner_depth
+            if "depth_refiner" in all_preds_data
+            else dt_coarse_refiner
+        )
+    else:
+        msg = f"{pred_name} extra data not in {all_preds_data.keys()}"
+        raise ValueError(msg)

happypose/pose_estimators/megapose/inference/pose_estimator.py

@@ -41,7 +41,10 @@
 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
-from happypose.toolbox.utils.tensor_collection import PandasTensorCollection
+from happypose.toolbox.utils.tensor_collection import (
+    PandasTensorCollection,
+    filter_top_pose_estimates,
+)
 from happypose.toolbox.utils.timer import Timer
 
 logger = get_logger(__name__)
@@ -585,11 +588,13 @@ def run_inference_pipeline(
             timing_str += f"coarse={coarse_extra_data['time']:.2f}, "
 
             # Extract top-K coarse hypotheses
-            data_TCO_filtered = self.filter_pose_estimates(
+            data_TCO_filtered = filter_top_pose_estimates(
                 data_TCO_coarse,
                 top_K=n_pose_hypotheses,
+                group_cols=["batch_im_id", "label", "instance_id"],
                 filter_field="coarse_logit",
             )
+
         else:
             data_TCO_coarse = coarse_estimates
             coarse_extra_data = None
@@ -614,9 +619,10 @@ def run_inference_pipeline(
         timing_str += f"scoring={scoring_extra_data['time']:.2f}, "
 
         # Extract the highest scoring pose estimate for each instance_id
-        data_TCO_final_scored = self.filter_pose_estimates(
+        data_TCO_final_scored = filter_top_pose_estimates(
             data_TCO_scored,
             top_K=1,
+            group_cols=["batch_im_id", "label", "instance_id"],
             filter_field="pose_logit",
         )
 
@@ -656,33 +662,3 @@ def run_inference_pipeline(
             extra_data["depth_refiner"] = {"preds": data_TCO_depth_refiner}
 
         return data_TCO_final, extra_data
-
-    def filter_pose_estimates(
-        self,
-        data_TCO: PoseEstimatesType,
-        top_K: int,
-        filter_field: str,
-        ascending: bool = False,
-    ) -> PoseEstimatesType:
-        """Filter the pose estimates by retaining only the top-K coarse model scores.
-
-        Retain only the top_K estimates corresponding to each hypothesis_id
-
-        Args:
-        ----
-            top_K: how many estimates to retain
-            filter_field: The field to filter estimates by
-        """
-        df = data_TCO.infos
-
-        group_cols = ["batch_im_id", "label", "instance_id"]
-        # Logic from https://stackoverflow.com/a/40629420
-        df = (
-            df.sort_values(filter_field, ascending=ascending)
-            .groupby(group_cols)
-            .head(top_K)
-        )
-
-        data_TCO_filtered = data_TCO[df.index.tolist()]
-
-        return data_TCO_filtered

happypose/pose_estimators/megapose/scripts/run_full_megapose_eval.py

@@ -50,17 +50,6 @@
 
 logger = get_logger(__name__)
 
-BOP_DATASET_NAMES = [
-    "lm",
-    "lmo",
-    "tless",
-    "tudl",
-    "icbin",
-    "itodd",
-    "hb",
-    "ycbv",
-    # 'hope',
-]
 
 BOP_TEST_DATASETS = [
     "lmo.bop19",
@@ -141,12 +130,12 @@ def run_full_eval(cfg: FullEvalConfig) -> None:
             if not cfg.skip_inference:
                 eval_out = run_eval(eval_cfg)
 
-            # If we are skpping the inference mimic the output that run_eval
+            # If we are skipping the inference, mimic the output that run_eval
             # would have produced so that we can run the bop_eval
             else:  # Otherwise hack the output so we can run the BOP eval
                 if get_rank() == 0:
                     results_dir = get_save_dir(eval_cfg)
-                    pred_keys = ["refiner/final"]
+                    pred_keys = ["coarse", "refiner/final"]
                     if eval_cfg.inference.run_depth_refiner:
                         pred_keys.append("depth_refiner")
                     eval_out = {
@@ -162,6 +151,11 @@ def run_full_eval(cfg: FullEvalConfig) -> None:
             # Run the bop eval for each type of prediction
             if cfg.run_bop_eval and get_rank() == 0:
                 bop_eval_keys = {"refiner/final", "depth_refiner"}
+                if cfg.eval_coarse_also:
+                    bop_eval_keys.add("coarse")
+
+                # Remove from evaluation predictions that were not produced at inference
+                # time
                 bop_eval_keys = bop_eval_keys.intersection(set(eval_out["pred_keys"]))
 
                 for method in bop_eval_keys:
@@ -174,7 +168,7 @@ def run_full_eval(cfg: FullEvalConfig) -> None:
                         split="test",
                         eval_dir=eval_out["save_dir"] / "bop_evaluation",
                         method=method,
-                        convert_only=False,
+                        convert_only=eval_cfg.convert_only,
                     )
                     bop_eval_cfgs.append(bop_eval_cfg)
 

happypose/toolbox/utils/tensor_collection.py

@@ -196,3 +196,35 @@ def __setstate__(self, state):
         self.__init__(state["infos"], **state["tensors"])
         self.meta = state["meta"]
         return
+
+
+def filter_top_pose_estimates(
+    data_TCO: PandasTensorCollection,
+    top_K: int,
+    group_cols: list[str],
+    filter_field: str,
+    ascending: bool = False,
+) -> PandasTensorCollection:
+    """Filter the pose estimates by retaining only the top-K coarse model scores.
+
+    Retain only the top_K estimates corresponding to each hypothesis_id
+
+    Args:
+        top_K: how many estimates to retain
+        group_cols: group of columns among which sorting should be done
+        filter_field: the field to filter estimates by
+        ascending: should filter_field
+    """
+
+    df = data_TCO.infos
+
+    # Logic from https://stackoverflow.com/a/40629420
+    df = (
+        df.sort_values(filter_field, ascending=ascending)
+        .groupby(group_cols)
+        .head(top_K)
+    )
+
+    data_TCO_filtered = data_TCO[df.index.tolist()]
+
+    return data_TCO_filtered

meh.sh

@@ -0,0 +1,22 @@
+#!/bin/bash -eux
+
+export HOME=/home/gsaurel
+export PATH=$HOME/.local/bin:$PATH
+export POETRY_VIRTUALENVS_IN_PROJECT=true
+
+rm -rf .venv local_data
+poetry env use /usr/bin/python
+poetry install --with dev -E cpu -E render -E evaluation
+
+mkdir local_data
+
+poetry run python -m happypose.toolbox.utils.download \
+    --megapose_models \
+    --examples \
+        crackers_example \
+    --cosypose_models  \
+        detector-bop-ycbv-pbr--970850 \
+        coarse-bop-ycbv-pbr--724183 \
+        refiner-bop-ycbv-pbr--604090
+
+poetry run python -m unittest