diff --git a/happypose/pose_estimators/megapose/evaluation/eval_config.py b/happypose/pose_estimators/megapose/evaluation/eval_config.py index 3a81d58e..398e52f3 100644 --- a/happypose/pose_estimators/megapose/evaluation/eval_config.py +++ b/happypose/pose_estimators/megapose/evaluation/eval_config.py @@ -17,7 +17,7 @@ # Standard Library from dataclasses import dataclass -from typing import List, Optional +from typing import Optional # MegaPose from happypose.pose_estimators.megapose.inference.types import InferenceConfig @@ -48,7 +48,8 @@ class EvalConfig: 2. If `run_id` is None, then use `config_id`, `run_comment`and `run_postfix` to create a `run_id` - In 2., the parameters of the config are set-up using the function `update_cfg_with_config_id`. + In 2., the parameters of the config are set-up using the function + `update_cfg_with_config_id`. """ # Network @@ -85,8 +86,8 @@ class EvalConfig: @dataclass class FullEvalConfig(EvalConfig): # Full eval - detection_coarse_types: Optional[List] = None - ds_names: Optional[List[str]] = None + detection_coarse_types: Optional[list] = None + ds_names: Optional[list[str]] = None run_bop_eval: bool = True eval_coarse_also: bool = False convert_only: bool = False diff --git a/happypose/pose_estimators/megapose/evaluation/prediction_runner.py b/happypose/pose_estimators/megapose/evaluation/prediction_runner.py index cd0a9bf3..eba04e3a 100644 --- a/happypose/pose_estimators/megapose/evaluation/prediction_runner.py +++ b/happypose/pose_estimators/megapose/evaluation/prediction_runner.py @@ -16,51 +16,35 @@ # Standard Library -import time from collections import defaultdict -from typing import Dict, Optional -from pathlib import Path - - -# Third Party -import numpy as np -import torch -from torch.utils.data import DataLoader -from tqdm import tqdm +from typing import Optional # MegaPose import happypose.pose_estimators.megapose import happypose.toolbox.utils.tensor_collection as tc -from happypose.pose_estimators.megapose.inference.pose_estimator import ( - PoseEstimator, + +# Third Party +import torch +from happypose.pose_estimators.megapose.evaluation.bop import ( + get_sam_detections, + load_sam_predictions, ) +from happypose.pose_estimators.megapose.inference.pose_estimator import PoseEstimator from happypose.pose_estimators.megapose.inference.types import ( DetectionsType, InferenceConfig, ObservationTensor, PoseEstimatesType, ) -from happypose.pose_estimators.megapose.config import BOP_DS_DIR -from happypose.pose_estimators.megapose.evaluation.bop import ( - get_sam_detections, - load_sam_predictions, -) - from happypose.pose_estimators.megapose.training.utils import CudaTimer from happypose.toolbox.datasets.samplers import DistributedSceneSampler -from happypose.toolbox.datasets.scene_dataset import ( - SceneDataset, - SceneObservation, - ObjectData, -) -from happypose.toolbox.utils.distributed import get_rank, get_tmp_dir, get_world_size -from happypose.toolbox.utils.logging import get_logger - +from happypose.toolbox.datasets.scene_dataset import SceneDataset, SceneObservation # Temporary -from happypose.toolbox.inference.utils import make_detections_from_object_data -import pandas as pd -import json +from happypose.toolbox.utils.distributed import get_rank, get_tmp_dir, get_world_size +from happypose.toolbox.utils.logging import get_logger +from torch.utils.data import DataLoader +from tqdm import tqdm logger = get_logger(__name__) @@ -81,7 +65,9 @@ def __init__( self.tmp_dir = get_tmp_dir() sampler = DistributedSceneSampler( - scene_ds, num_replicas=self.world_size, rank=self.rank + scene_ds, + num_replicas=self.world_size, + rank=self.rank, ) self.sampler = sampler self.scene_ds = scene_ds @@ -104,12 +90,13 @@ def run_inference_pipeline( gt_detections: DetectionsType, sam_detections: DetectionsType, initial_estimates: Optional[PoseEstimatesType] = None, - ) -> Dict[str, PoseEstimatesType]: + ) -> dict[str, PoseEstimatesType]: """Runs inference pipeline, extracts the results. Returns: A dict with keys - 'final': final preds - - 'refiner/final': preds at final refiner iteration (before depth refinement) + - 'refiner/final': preds at final refiner iteration (before depth + refinement) - 'depth_refinement': preds after depth refinement. @@ -128,16 +115,15 @@ def run_inference_pipeline( run_detector = True else: - raise ValueError( - f"Unknown detection type {self.inference_cfg.detection_type}" - ) + msg = f"Unknown detection type {self.inference_cfg.detection_type}" + raise ValueError(msg) coarse_estimates = None if self.inference_cfg.coarse_estimation_type == "external": # TODO (ylabbe): This is hacky, clean this for modelnet eval. coarse_estimates = initial_estimates coarse_estimates = happypose.toolbox.inference.utils.add_instance_id( - coarse_estimates + coarse_estimates, ) coarse_estimates.infos["instance_id"] = 0 run_detector = False @@ -159,13 +145,13 @@ def run_inference_pipeline( # - 'refiner/iteration=5` # - `depth_refiner` # Note: Since we support multi-hypotheses we need to potentially - # 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. + # 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. + ref_str = f"refiner/iteration={self.inference_cfg.n_refiner_iterations}" all_preds = { "final": preds, - f"refiner/iteration={self.inference_cfg.n_refiner_iterations}": extra_data[ - "refiner" - ]["preds"], + ref_str: extra_data["refiner"]["preds"], "refiner/final": extra_data["refiner"]["preds"], "coarse": extra_data["coarse"]["preds"], "coarse_filter": extra_data["coarse_filter"]["preds"], @@ -183,17 +169,18 @@ def run_inference_pipeline( 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(): + for _k, v in all_preds.items(): if "mask" in v.tensors: - breakpoint() + # breakpoint() v.delete_tensor("mask") return all_preds, all_preds_data def get_predictions( - self, pose_estimator: PoseEstimator - ) -> Dict[str, PoseEstimatesType]: - """Runs predictions + self, + pose_estimator: PoseEstimator, + ) -> dict[str, PoseEstimatesType]: + """Runs predictions. Returns: A dict with keys - 'refiner/iteration=1` @@ -204,7 +191,6 @@ def get_predictions( """ - predictions_list = defaultdict(list) ###### @@ -214,7 +200,8 @@ def get_predictions( # Temporary solution if self.inference_cfg.detection_type == "sam": df_all_dets, df_targets = load_sam_predictions( - self.scene_ds.ds_dir.name, self.scene_ds.ds_dir + self.scene_ds.ds_dir.name, + self.scene_ds.ds_dir, ) for n, data in enumerate(tqdm(self.dataloader)): @@ -274,19 +261,20 @@ 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 pred_name, pred in all_preds.items(): pred.infos["time"] = dt_det + compute_pose_est_total_time( - all_preds_data, pred_name + 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 = dict() + predictions = {} for k, v in predictions_list.items(): predictions[k] = tc.concatenate(v) @@ -315,4 +303,5 @@ def compute_pose_est_total_time(all_preds_data: dict, pred_name: str): else dt_coarse_refiner ) else: - raise ValueError(f"{pred_name} extra data not in {all_preds_data.keys()}") + 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 5018901a..853ab633 100644 --- a/happypose/pose_estimators/megapose/inference/pose_estimator.py +++ b/happypose/pose_estimators/megapose/inference/pose_estimator.py @@ -1,5 +1,4 @@ -""" -Copyright (c) 2022 Inria & NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""Copyright (c) 2022 Inria & NVIDIA CORPORATION & AFFILIATES. All rights reserved. Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. @@ -20,8 +19,7 @@ # Standard Library import time from collections import defaultdict -from dataclasses import dataclass -from typing import Any, Optional, Tuple +from typing import Any # Third Party import numpy as np @@ -30,16 +28,10 @@ from torch.utils.data import DataLoader, TensorDataset # MegaPose -import happypose.pose_estimators.megapose as megapose import happypose.toolbox.inference.utils import happypose.toolbox.utils.tensor_collection as tc -from happypose.pose_estimators.megapose.inference.depth_refiner import ( - DepthRefiner, -) -from happypose.pose_estimators.megapose.training.utils import ( - CudaTimer, - SimpleTimer, -) +from happypose.pose_estimators.megapose.inference.depth_refiner import DepthRefiner +from happypose.pose_estimators.megapose.training.utils import CudaTimer, SimpleTimer from happypose.toolbox.inference.pose_estimator import PoseEstimationModule from happypose.toolbox.inference.types import ( DetectionsType, @@ -65,10 +57,10 @@ class PoseEstimator(PoseEstimationModule): def __init__( self, - refiner_model: Optional[torch.nn.Module] = None, - coarse_model: Optional[torch.nn.Module] = None, - detector_model: Optional[torch.nn.Module] = None, - depth_refiner: Optional[DepthRefiner] = None, + refiner_model: torch.nn.Module | None = None, + coarse_model: torch.nn.Module | None = None, + detector_model: torch.nn.Module | None = None, + depth_refiner: DepthRefiner | None = None, bsz_objects: int = 8, bsz_images: int = 256, SO3_grid_size: int = 576, @@ -93,9 +85,8 @@ def __init__( self.cfg = self.coarse_model.cfg self.mesh_db = self.coarse_model.mesh_db else: - raise ValueError( - "At least one of refiner_model or " " coarse_model must be specified." - ) + msg = "At least one of refiner_model or coarse_model must be specified." + raise ValueError(msg) self.eval() @@ -103,7 +94,7 @@ def __init__( self.keep_all_coarse_outputs = False self.refiner_outputs = None self.coarse_outputs = None - self.debug_dict: dict = dict() + self.debug_dict: dict = {} def load_SO3_grid(self, grid_size: int) -> None: """Loads the SO(3) grid.""" @@ -119,14 +110,14 @@ def forward_refiner( keep_all_outputs: bool = False, cuda_timer: bool = False, **refiner_kwargs, - ) -> Tuple[dict, dict]: + ) -> tuple[dict, dict]: """Runs the refiner model for the specified number of iterations. - Will actually use the batched_model_predictions to stay within batch size limit. - Returns: + Returns + ------- (preds, extra_data) preds: @@ -139,7 +130,6 @@ def forward_refiner( A dict containing additional information such as timing """ - timer = Timer() timer.start() @@ -224,7 +214,7 @@ def forward_refiner( logger.debug( f"Pose prediction on {B} poses (n_iterations={n_iterations}):" - f" {timer.stop()}" + f" {timer.stop()}", ) return preds, extra_data @@ -236,15 +226,13 @@ def forward_scoring_model( data_TCO: PoseEstimatesType, cuda_timer: bool = False, return_debug_data: bool = False, - ) -> Tuple[PoseEstimatesType, dict]: + ) -> tuple[PoseEstimatesType, dict]: """Score the estimates using the coarse model. - Adds the 'pose_score' field to data_TCO.infos Modifies PandasTensorCollection in-place. """ - start_time = time.time() assert self.coarse_model is not None @@ -295,7 +283,7 @@ def forward_scoring_model( images_crop_list.append(out_["images_crop"]) renders_list.append(out_["renders"]) - debug_data = dict() + debug_data = {} # Combine together the data from the different batches logits = torch.cat(logits_list) @@ -304,8 +292,8 @@ def forward_scoring_model( images_crop: torch.tensor = torch.cat(images_crop_list) renders: torch.tensor = torch.cat(renders_list) - H = images_crop.shape[2] - W = images_crop.shape[3] + images_crop.shape[2] + images_crop.shape[3] debug_data = { "images_crop": images_crop, @@ -317,7 +305,10 @@ def forward_scoring_model( elapsed = time.time() - start_time - timing_str = f"time: {elapsed:.2f}, model_time: {model_time:.2f}, render_time: {render_time:.2f}" + timing_str = ( + f"time: {elapsed:.2f}, model_time: {model_time:.2f}, " + f"render_time: {render_time:.2f}" + ) extra_data = { "render_time": render_time, @@ -340,13 +331,12 @@ def forward_coarse_model( detections: DetectionsType, cuda_timer: bool = False, return_debug_data: bool = False, - ) -> Tuple[PoseEstimatesType, dict]: + ) -> tuple[PoseEstimatesType, dict]: """Generates pose hypotheses and scores them with the coarse model. - Generates coarse hypotheses using the SO(3) grid. - Scores them using the coarse model. """ - start_time = time.time() happypose.toolbox.inference.types.assert_detections_valid(detections) @@ -455,7 +445,7 @@ def forward_coarse_model( TCO = torch.cat(TCO_init) TCO_reshape = TCO.reshape([B, M, 4, 4]) - debug_data = dict() + debug_data = {} if return_debug_data: images_crop = torch.cat(images_crop_list) @@ -474,7 +464,10 @@ def forward_coarse_model( elapsed = time.time() - start_time - timing_str = f"time: {elapsed:.2f}, model_time: {model_time:.2f}, render_time: {render_time:.2f}" + timing_str = ( + f"time: {elapsed:.2f}, model_time: {model_time:.2f}, " + f"render_time: {render_time:.2f}" + ) extra_data = { "render_time": render_time, @@ -499,21 +492,22 @@ def forward_detection_model( **kwargs: Any, ) -> DetectionsType: """Runs the detector.""" - return self.detector_model.get_detections(observation, *args, **kwargs) def run_depth_refiner( self, observation: ObservationTensor, predictions: PoseEstimatesType, - ) -> Tuple[PoseEstimatesType, dict]: + ) -> tuple[PoseEstimatesType, dict]: """Runs the depth refiner.""" assert self.depth_refiner is not None, "You must specify a depth refiner" depth = observation.depth K = observation.K refined_preds, extra_data = self.depth_refiner.refine_poses( - predictions, depth=depth, K=K + predictions, + depth=depth, + K=K, ) return refined_preds, extra_data @@ -522,18 +516,18 @@ def run_depth_refiner( def run_inference_pipeline( self, observation: ObservationTensor, - detections: Optional[DetectionsType] = None, - run_detector: Optional[bool] = None, + detections: DetectionsType | None = None, + run_detector: bool | None = None, n_refiner_iterations: int = 5, n_pose_hypotheses: int = 1, keep_all_refiner_outputs: bool = False, - detection_filter_kwargs: Optional[dict] = None, + detection_filter_kwargs: dict | None = None, run_depth_refiner: bool = False, - bsz_images: Optional[int] = None, - bsz_objects: Optional[int] = None, + bsz_images: int | None = None, + bsz_objects: int | None = None, cuda_timer: bool = False, - coarse_estimates: Optional[PoseEstimatesType] = None, - ) -> Tuple[PoseEstimatesType, dict]: + coarse_estimates: PoseEstimatesType | None = None, + ) -> tuple[PoseEstimatesType, dict]: """Runs the entire pose estimation pipeline. Performs the following steps @@ -545,13 +539,13 @@ def run_inference_pipeline( 5. Score refined hypotheses 6. Select highest scoring refined hypotheses. - Returns: + Returns + ------- data_TCO_final: final predictions data: Dict containing additional data about the different steps in the pipeline. """ - timing_str = "" timer = SimpleTimer() timer.start() @@ -581,7 +575,8 @@ def run_inference_pipeline( # Filter detections if detection_filter_kwargs is not None: detections = happypose.toolbox.inference.utils.filter_detections( - detections, **detection_filter_kwargs + detections, + **detection_filter_kwargs, ) # Run the coarse estimator using detections @@ -635,7 +630,8 @@ def run_inference_pipeline( if run_depth_refiner: depth_refiner_start = time.time() data_TCO_depth_refiner, _ = self.run_depth_refiner( - observation, data_TCO_final_scored + observation, + data_TCO_final_scored, ) data_TCO_final = data_TCO_depth_refiner depth_refiner_time = time.time() - depth_refiner_start @@ -647,7 +643,7 @@ def run_inference_pipeline( timer.stop() timing_str = f"total={timer.elapsed():.2f}, {timing_str}" - extra_data: dict = dict() + extra_data: dict = {} extra_data["coarse"] = {"preds": data_TCO_coarse, "data": coarse_extra_data} extra_data["coarse_filter"] = {"preds": data_TCO_filtered} extra_data["refiner_all_hypotheses"] = { 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 42f1c7d5..e2ce6cb0 100644 --- a/happypose/pose_estimators/megapose/scripts/run_full_megapose_eval.py +++ b/happypose/pose_estimators/megapose/scripts/run_full_megapose_eval.py @@ -1,5 +1,4 @@ -""" -Copyright (c) 2022 Inria & NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""Copyright (c) 2022 Inria & NVIDIA CORPORATION & AFFILIATES. All rights reserved. Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. @@ -19,39 +18,26 @@ import copy import os from pathlib import Path -from typing import Dict, Optional, Tuple - -# Third Party -from omegaconf import OmegaConf # MegaPose -from happypose.pose_estimators.megapose.bop_config import ( - PBR_COARSE, - PBR_DETECTORS, - PBR_REFINER, - SYNT_REAL_COARSE, - SYNT_REAL_DETECTORS, - SYNT_REAL_REFINER, -) +from happypose.pose_estimators.megapose.bop_config import PBR_DETECTORS from happypose.pose_estimators.megapose.config import ( DEBUG_RESULTS_DIR, - EXP_DIR, MODELNET_TEST_CATEGORIES, RESULTS_DIR, ) +from happypose.pose_estimators.megapose.evaluation.bop import run_evaluation from happypose.pose_estimators.megapose.evaluation.eval_config import ( BOPEvalConfig, EvalConfig, FullEvalConfig, HardwareConfig, ) - from happypose.pose_estimators.megapose.evaluation.evaluation import ( - get_save_dir, generate_save_key, + get_save_dir, run_eval, ) -from happypose.pose_estimators.megapose.evaluation.bop import run_evaluation from happypose.toolbox.utils.distributed import ( get_rank, get_world_size, @@ -59,6 +45,9 @@ ) from happypose.toolbox.utils.logging import get_logger, set_logging_level +# Third Party +from omegaconf import OmegaConf + logger = get_logger(__name__) @@ -83,7 +72,7 @@ def create_eval_cfg( detection_type: str, coarse_estimation_type: str, ds_name: str, -) -> Tuple[str, EvalConfig]: +) -> tuple[str, EvalConfig]: cfg = copy.deepcopy(cfg) cfg.inference.detection_type = detection_type @@ -101,7 +90,8 @@ def create_eval_cfg( elif detection_type == "sam": pass else: - raise ValueError(f"Unknown detector type {cfg.detector_type}") + msg = f"Unknown detector type {cfg.detector_type}" + raise ValueError(msg) name = generate_save_key(detection_type, coarse_estimation_type) @@ -123,16 +113,19 @@ def run_full_eval(cfg: FullEvalConfig) -> None: # Iterate over each dataset for ds_name in cfg.ds_names: # create the EvalConfig objects that we will call `run_eval` on - eval_configs: Dict[str, EvalConfig] = dict() + eval_configs: dict[str, EvalConfig] = {} for detection_type, coarse_estimation_type in cfg.detection_coarse_types: name, cfg_ = create_eval_cfg( - cfg, detection_type, coarse_estimation_type, ds_name + cfg, + detection_type, + coarse_estimation_type, + ds_name, ) eval_configs[name] = cfg_ # For each eval_cfg run the evaluation. # Note that the results get saved to disk - for save_key, eval_cfg in eval_configs.items(): + for _save_key, eval_cfg in eval_configs.items(): # Run the inference if not cfg.skip_inference: eval_out = run_eval(eval_cfg) @@ -152,12 +145,12 @@ def run_full_eval(cfg: FullEvalConfig) -> None: } assert Path( - eval_out["results_path"] + eval_out["results_path"], ).is_file(), f"The file {eval_out['results_path']} doesn't exist" # Run the bop eval for each type of prediction if cfg.run_bop_eval and get_rank() == 0: - bop_eval_keys = set(("refiner/final", "depth_refiner")) + bop_eval_keys = {"refiner/final", "depth_refiner"} if cfg.eval_coarse_also: bop_eval_keys.add("coarse") @@ -165,7 +158,7 @@ def run_full_eval(cfg: FullEvalConfig) -> None: bop_eval_keys = bop_eval_keys.intersection(set(eval_out["pred_keys"])) for method in bop_eval_keys: - if not "bop19" in ds_name: + if "bop19" not in ds_name: continue bop_eval_cfg = BOPEvalConfig(