diff --git a/happypose/pose_estimators/megapose/evaluation/bop.py b/happypose/pose_estimators/megapose/evaluation/bop.py index c879ad86..519e9429 100644 --- a/happypose/pose_estimators/megapose/evaluation/bop.py +++ b/happypose/pose_estimators/megapose/evaluation/bop.py @@ -36,7 +36,7 @@ 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 +from happypose.toolbox.utils.tensor_collection import PandasTensorCollection, filter_top_pose_estimates device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') @@ -142,39 +142,9 @@ def convert_results_to_coco(results_path, out_json_path, detection_method): return - - -def filter_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 - filter_field: The field to filter estimates by - """ - - # TODO: refactor with definition in pose_estimator.py - - 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 - - def get_best_coarse_predictions(coarse_preds: PandasTensorCollection): group_cols = ["scene_id", "view_id", "label", "instance_id"] - coarse_preds = filter_pose_estimates(coarse_preds, top_K=1, group_cols=group_cols, filter_field='coarse_score', ascending=False) + 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 diff --git a/happypose/pose_estimators/megapose/inference/pose_estimator.py b/happypose/pose_estimators/megapose/inference/pose_estimator.py index e54af605..8a242c17 100644 --- a/happypose/pose_estimators/megapose/inference/pose_estimator.py +++ b/happypose/pose_estimators/megapose/inference/pose_estimator.py @@ -49,7 +49,7 @@ 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__) @@ -591,9 +591,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_coarse, top_K=n_pose_hypotheses, filter_field="coarse_logit" + 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 @@ -618,8 +622,11 @@ 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_scored, top_K=1, filter_field="pose_logit" + data_TCO_final_scored = self.filter_top_pose_estimates( + data_TCO_scored, + top_K=1, + group_cols=["batch_im_id", "label", "instance_id"], + filter_field="pose_logit" ) # Optionally run ICP or TEASER++ @@ -649,29 +656,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 \ No newline at end of file diff --git a/happypose/toolbox/utils/tensor_collection.py b/happypose/toolbox/utils/tensor_collection.py index f0c35c83..2c005d23 100644 --- a/happypose/toolbox/utils/tensor_collection.py +++ b/happypose/toolbox/utils/tensor_collection.py @@ -194,3 +194,32 @@ 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 diff --git a/notebooks/megapose/megapose_estimator_visualization.ipynb b/notebooks/megapose/megapose_estimator_visualization.ipynb index 9a557e5e..be23a8b5 100644 --- a/notebooks/megapose/megapose_estimator_visualization.ipynb +++ b/notebooks/megapose/megapose_estimator_visualization.ipynb @@ -243,6 +243,8 @@ "metadata": {}, "outputs": [], "source": [ + "from happypose.toolbox.utils.tensor_collection import filter_pose_estimates\n", + "\n", "# Options for inference\n", "use_gt_detections = True # Note, if you aren't using gt_detections then this should be false\n", "n_refiner_iterations = 5\n", @@ -294,10 +296,11 @@ " f\"model_time={extra_data['model_time']:.2f}, render_time={extra_data['render_time']:.2f}\")\n", " \n", " # Extract top-K coarse hypotheses\n", - " data_TCO_filtered = pose_estimator.filter_pose_estimates(data_TCO_coarse, \n", - " top_K=n_pose_hypotheses, \n", - " filter_field='coarse_logit')\n", - " \n", + " data_TCO_filtered = filter_pose_estimates(data_TCO_coarse,\n", + " top_K=n_pose_hypotheses, \n", + " group_cols=[\"batch_im_id\", \"label\", \"instance_id\"], \n", + " filter_field='coarse_logit')\n", + "\n", " # Refine the top_K coarse hypotheses\n", " preds, extra_data = pose_estimator.forward_refiner(observation_tensor, data_TCO_filtered, \n", " n_iterations=n_refiner_iterations, keep_all_outputs=True)\n", @@ -311,7 +314,10 @@ " data_TCO_scored, extra_data = pose_estimator.forward_scoring_model(observation_tensor, data_TCO_refined)\n", "\n", " # Extract the highest scoring pose estimate for each instance_id\n", - " data_TCO_final = pose_estimator.filter_pose_estimates(data_TCO_scored, top_K=1, filter_field='pose_logit')\n", + " data_TCO_final = filter_pose_estimates(data_TCO_scored, \n", + " top_K=1, \n", + " group_cols=[\"batch_im_id\", \"label\", \"instance_id\"], \n", + " filter_field='pose_logit')\n", " \n", " \n", " if run_depth_refiner:\n", @@ -969,4 +975,4 @@ }, "nbformat": 4, "nbformat_minor": 5 -} \ No newline at end of file +}