diff --git a/happypose/pose_estimators/megapose/evaluation/bop.py b/happypose/pose_estimators/megapose/evaluation/bop.py index 57bf0dc9..733349e1 100644 --- a/happypose/pose_estimators/megapose/evaluation/bop.py +++ b/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)) diff --git a/happypose/pose_estimators/megapose/evaluation/eval_config.py b/happypose/pose_estimators/megapose/evaluation/eval_config.py index 1eccf2d8..9aba8b1e 100644 --- a/happypose/pose_estimators/megapose/evaluation/eval_config.py +++ b/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 diff --git a/happypose/pose_estimators/megapose/evaluation/prediction_runner.py b/happypose/pose_estimators/megapose/evaluation/prediction_runner.py index f722f2d6..2f8a3ebf 100644 --- a/happypose/pose_estimators/megapose/evaluation/prediction_runner.py +++ b/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,14 +260,17 @@ 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 = {} @@ -271,3 +278,32 @@ def get_predictions( 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) diff --git a/happypose/pose_estimators/megapose/inference/pose_estimator.py b/happypose/pose_estimators/megapose/inference/pose_estimator.py index 19238834..853ab633 100644 --- a/happypose/pose_estimators/megapose/inference/pose_estimator.py +++ b/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 diff --git a/happypose/pose_estimators/megapose/scripts/run_full_megapose_eval.py b/happypose/pose_estimators/megapose/scripts/run_full_megapose_eval.py index 2059e1ad..de450e5b 100644 --- a/happypose/pose_estimators/megapose/scripts/run_full_megapose_eval.py +++ b/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) diff --git a/happypose/toolbox/utils/tensor_collection.py b/happypose/toolbox/utils/tensor_collection.py index 63874097..c1bd696f 100644 --- a/happypose/toolbox/utils/tensor_collection.py +++ b/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 diff --git a/meh.sh b/meh.sh new file mode 100755 index 00000000..d60db239 --- /dev/null +++ b/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 diff --git a/notebooks/megapose/megapose_estimator_visualization.ipynb b/notebooks/megapose/megapose_estimator_visualization.ipynb index de04b7bd..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",