diff --git a/happypose/pose_estimators/cosypose/cosypose/datasets/augmentations.py b/happypose/pose_estimators/cosypose/cosypose/datasets/augmentations.py index 362b7b52..035b2c4d 100644 --- a/happypose/pose_estimators/cosypose/cosypose/datasets/augmentations.py +++ b/happypose/pose_estimators/cosypose/cosypose/datasets/augmentations.py @@ -143,7 +143,6 @@ def __call__(self, im, mask, obs): class VOCBackgroundAugmentation(BackgroundAugmentation): def __init__(self, voc_root, p=0.3): - print("voc_root =", voc_root) image_dataset = VOCSegmentation( root=voc_root, year="2012", diff --git a/happypose/pose_estimators/cosypose/cosypose/evaluation/meters/utils.py b/happypose/pose_estimators/cosypose/cosypose/evaluation/meters/utils.py index 18fc197e..8847d330 100644 --- a/happypose/pose_estimators/cosypose/cosypose/evaluation/meters/utils.py +++ b/happypose/pose_estimators/cosypose/cosypose/evaluation/meters/utils.py @@ -25,7 +25,7 @@ def add_inst_num( group_keys=["scene_id", "view_id", "label"], key="pred_inst_num", ): - inst_num = np.empty(len(infos), dtype=np.int) + inst_num = np.empty(len(infos), dtype=int) for _group_name, group_ids in infos.groupby(group_keys).groups.items(): inst_num[group_ids.values] = np.arange(len(group_ids)) infos[key] = inst_num diff --git a/happypose/pose_estimators/cosypose/cosypose/evaluation/prediction_runner.py b/happypose/pose_estimators/cosypose/cosypose/evaluation/prediction_runner.py index 33ec1f24..f8ba10df 100644 --- a/happypose/pose_estimators/cosypose/cosypose/evaluation/prediction_runner.py +++ b/happypose/pose_estimators/cosypose/cosypose/evaluation/prediction_runner.py @@ -95,18 +95,18 @@ def run_inference_pipeline( """ - if self.inference_cfg.detection_type == "gt": + if self.inference_cfg['detection_type'] == "gt": detections = gt_detections run_detector = False - elif self.inference_cfg.detection_type == "detector": + elif self.inference_cfg['detection_type'] == "detector": detections = None run_detector = True else: - msg = 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": + 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( @@ -137,7 +137,7 @@ def run_inference_pipeline( all_preds = {} data_TCO_refiner = extra_data["refiner"]["preds"] - k_0 = f"refiner/iteration={self.inference_cfg.n_refiner_iterations}" + k_0 = f"refiner/iteration={self.inference_cfg['n_refiner_iterations']}" all_preds = { "final": preds, k_0: data_TCO_refiner, @@ -145,7 +145,7 @@ def run_inference_pipeline( "coarse": extra_data["coarse"]["preds"], } - if self.inference_cfg.run_depth_refiner: + if self.inference_cfg['run_depth_refiner']: all_preds["depth_refiner"] = extra_data["depth_refiner"]["preds"] # Remove any mask tensors @@ -174,43 +174,46 @@ def get_predictions( """ predictions_list = defaultdict(list) for n, data in enumerate(tqdm(self.dataloader)): - # data is a dict - rgb = data["rgb"] - depth = None - K = data["cameras"].K - gt_detections = data["gt_detections"].cuda() - - initial_data = None - if data["initial_data"]: - initial_data = data["initial_data"].cuda() - - obs_tensor = ObservationTensor.from_torch_batched(rgb, depth, K) - obs_tensor = obs_tensor.cuda() - - # GPU warmup for timing - if n == 0: + if n < 3: + # data is a dict + rgb = data["rgb"] + depth = None + K = data["cameras"].K + gt_detections = data["gt_detections"].cuda() + + initial_data = None + if data["initial_data"]: + initial_data = data["initial_data"].cuda() + + obs_tensor = ObservationTensor.from_torch_batched(rgb, depth, K) + obs_tensor = obs_tensor.cuda() + + # GPU warmup for timing + if n == 0: + with torch.no_grad(): + self.run_inference_pipeline( + pose_estimator, + obs_tensor, + gt_detections, + initial_estimates=initial_data, + ) + + cuda_timer = CudaTimer() + cuda_timer.start() with torch.no_grad(): - self.run_inference_pipeline( + all_preds = self.run_inference_pipeline( pose_estimator, obs_tensor, gt_detections, initial_estimates=initial_data, ) + cuda_timer.end() + cuda_timer.elapsed() - cuda_timer = CudaTimer() - cuda_timer.start() - with torch.no_grad(): - all_preds = self.run_inference_pipeline( - pose_estimator, - obs_tensor, - gt_detections, - initial_estimates=initial_data, - ) - cuda_timer.end() - cuda_timer.elapsed() - - for k, v in all_preds.items(): - predictions_list[k].append(v) + for k, v in all_preds.items(): + predictions_list[k].append(v) + else: + break # Concatenate the lists of PandasTensorCollections predictions = {} diff --git a/happypose/pose_estimators/cosypose/cosypose/evaluation/runner_utils.py b/happypose/pose_estimators/cosypose/cosypose/evaluation/runner_utils.py index fb48bfff..b360431c 100644 --- a/happypose/pose_estimators/cosypose/cosypose/evaluation/runner_utils.py +++ b/happypose/pose_estimators/cosypose/cosypose/evaluation/runner_utils.py @@ -14,8 +14,8 @@ def run_pred_eval(pred_runner, pred_kwargs, eval_runner, eval_preds=None): all_predictions = {} for pred_prefix, pred_kwargs_n in pred_kwargs.items(): - print("Prediction :", pred_prefix) - preds = pred_runner.get_predictions(**pred_kwargs_n) + pose_predictor = pred_kwargs_n['pose_predictor'] + preds = pred_runner.get_predictions(pose_predictor) for preds_name, preds_n in preds.items(): all_predictions[f"{pred_prefix}/{preds_name}"] = preds_n diff --git a/happypose/pose_estimators/cosypose/cosypose/models/pose.py b/happypose/pose_estimators/cosypose/cosypose/models/pose.py index 48e2a5c3..95ee2ede 100644 --- a/happypose/pose_estimators/cosypose/cosypose/models/pose.py +++ b/happypose/pose_estimators/cosypose/cosypose/models/pose.py @@ -167,6 +167,7 @@ def forward(self, images, K, labels, TCO, n_iterations=1): K_crop=K_crop, boxes_rend=boxes_rend, boxes_crop=boxes_crop, + model_outputs=model_outputs ) TCO_input = TCO_output diff --git a/happypose/pose_estimators/cosypose/cosypose/scripts/run_pose_training.py b/happypose/pose_estimators/cosypose/cosypose/scripts/run_pose_training.py index 44551e50..b9f52d2a 100644 --- a/happypose/pose_estimators/cosypose/cosypose/scripts/run_pose_training.py +++ b/happypose/pose_estimators/cosypose/cosypose/scripts/run_pose_training.py @@ -10,6 +10,9 @@ logger = get_logger(__name__) +import warnings +warnings.filterwarnings("ignore") + def make_cfg(args): cfg = argparse.ArgumentParser("").parse_args([]) if args.config: @@ -56,7 +59,7 @@ def make_cfg(args): cfg.n_pose_dims = 9 cfg.n_rendering_workers = N_WORKERS cfg.refiner_run_id_for_test = None - cfg.coarse_run_id_for_test = None + cfg.coarse_run_id_for_test = "coarse-bop-ycbv-pbr--724183" # Optimizer cfg.lr = 3e-4 @@ -66,7 +69,7 @@ def make_cfg(args): cfg.clip_grad_norm = 0.5 # Training - cfg.batch_size = 32 + cfg.batch_size = 16 cfg.epoch_size = 115200 cfg.n_epochs = 700 cfg.n_dataloader_workers = N_WORKERS diff --git a/happypose/pose_estimators/cosypose/cosypose/training/pose_forward_loss.py b/happypose/pose_estimators/cosypose/cosypose/training/pose_forward_loss.py index 4cb2d740..17c8e6b9 100644 --- a/happypose/pose_estimators/cosypose/cosypose/training/pose_forward_loss.py +++ b/happypose/pose_estimators/cosypose/cosypose/training/pose_forward_loss.py @@ -63,10 +63,11 @@ def h_pose(model, mesh_db, data, meters, cfg, n_iterations=1, input_generator="f losses_TCO_iter = [] for n in range(n_iterations): iter_outputs = outputs[f"iteration={n+1}"] - K_crop = iter_outputs["K_crop"] - TCO_input = iter_outputs["TCO_input"] - TCO_pred = iter_outputs["TCO_output"] - model_outputs = iter_outputs["model_outputs"] + K_crop = iter_outputs.K_crop + TCO_input = iter_outputs.TCO_input + TCO_pred = iter_outputs.TCO_output + model_outputs = iter_outputs.model_outputs + if cfg.loss_disentangled: if cfg.n_pose_dims == 9: diff --git a/happypose/pose_estimators/cosypose/cosypose/training/train_pose.py b/happypose/pose_estimators/cosypose/cosypose/training/train_pose.py index 01040496..9d57094c 100644 --- a/happypose/pose_estimators/cosypose/cosypose/training/train_pose.py +++ b/happypose/pose_estimators/cosypose/cosypose/training/train_pose.py @@ -97,7 +97,7 @@ def load_model(run_id): if run_id is None: return None run_dir = EXP_DIR / run_id - cfg = yaml.load((run_dir / "config.yaml").read_text(), Loader=yaml.FullLoader) + cfg = yaml.load((run_dir / "config.yaml").read_text(), Loader=yaml.Loader) cfg = check_update_config(cfg) model = ( create_model_pose( @@ -422,7 +422,7 @@ def train_epoch(): iterator = tqdm(ds_iter_train, ncols=80) t = time.time() for n, sample in enumerate(iterator): - if n < 5: + if n < 3: if n > 0: meters_time["data"].add(time.time() - t) @@ -453,7 +453,7 @@ def train_epoch(): lr_scheduler_warmup.step() t = time.time() else: - continue + break if epoch >= args.n_epochs_warmup: lr_scheduler.step() @@ -461,11 +461,11 @@ def train_epoch(): def validation(): model.eval() for n, sample in enumerate(tqdm(ds_iter_val, ncols=80)): - if n < 5: + if n < 3: loss = h(data=sample, meters=meters_val) meters_val["loss_total"].add(loss.item()) else: - continue + break @torch.no_grad() def test(): diff --git a/happypose/pose_estimators/megapose/evaluation/runner_utils.py b/happypose/pose_estimators/megapose/evaluation/runner_utils.py index 2e90079d..4579b80c 100644 --- a/happypose/pose_estimators/megapose/evaluation/runner_utils.py +++ b/happypose/pose_estimators/megapose/evaluation/runner_utils.py @@ -30,7 +30,6 @@ def run_pred_eval(pred_runner, pred_kwargs, eval_runner, eval_preds=None): all_predictions = {} for pred_prefix, pred_kwargs_n in pred_kwargs.items(): - print("Prediction :", pred_prefix) preds = pred_runner.get_predictions(**pred_kwargs_n) for preds_name, preds_n in preds.items(): all_predictions[f"{pred_prefix}/{preds_name}"] = preds_n diff --git a/happypose/pose_estimators/megapose/models/pose_rigid.py b/happypose/pose_estimators/megapose/models/pose_rigid.py index f37de84e..094384ea 100644 --- a/happypose/pose_estimators/megapose/models/pose_rigid.py +++ b/happypose/pose_estimators/megapose/models/pose_rigid.py @@ -59,6 +59,7 @@ class PosePredictorOutputCosypose: K_crop: torch.Tensor boxes_rend: torch.Tensor boxes_crop: torch.Tensor + model_outputs: torch.Tensor @dataclass diff --git a/happypose/toolbox/lib3d/rigid_mesh_database.py b/happypose/toolbox/lib3d/rigid_mesh_database.py index 13263897..c24c919c 100644 --- a/happypose/toolbox/lib3d/rigid_mesh_database.py +++ b/happypose/toolbox/lib3d/rigid_mesh_database.py @@ -146,6 +146,8 @@ def n_sym_mapping(self): return {label: obj["n_sym"] for label, obj in self.infos.items()} def select(self, labels): + print("label to id =", self.label_to_id) + print("labels = ", labels) ids = [self.label_to_id[label] for label in labels] return Meshes( infos=[self.infos[label] for label in labels],