diff --git a/core/gdrn_modeling/engine/gdrn_evaluator.py b/core/gdrn_modeling/engine/gdrn_evaluator.py index e1cb3186cdb028d93a373a805b32b7c4a35d7252..a2bf2f40e59f7494c6eb3633cd6981f992fcef34 100644 --- a/core/gdrn_modeling/engine/gdrn_evaluator.py +++ b/core/gdrn_modeling/engine/gdrn_evaluator.py @@ -512,7 +512,7 @@ class GDRN_Evaluator(DatasetEvaluator): rot_est = out_rots[out_i] trans_est = out_transes[out_i] pose_est = np.hstack([rot_est, trans_est.reshape(3, 1)]) - depth_sensor_crop = _input['roi_depth'][inst_i].cpu().numpy().copy().squeeze() + depth_sensor_crop = cv2.resize(_input['roi_depth'][inst_i][-1].cpu().numpy().copy().squeeze(), (self.out_res, self.out_res)) depth_sensor_mask_crop = depth_sensor_crop > 0 net_cfg = cfg.MODEL.POSE_NET diff --git a/core/gdrn_modeling/models/GDRN.py b/core/gdrn_modeling/models/GDRN.py index 1a0e3becb919e052a30b36cacb7e843227cb8a0c..98cddad91c3e6060dcc5b8e24675c0c9ee36c7d3 100644 --- a/core/gdrn_modeling/models/GDRN.py +++ b/core/gdrn_modeling/models/GDRN.py @@ -198,7 +198,7 @@ class GDRN(nn.Module): if not do_loss: # test out_dict = {"rot": pred_ego_rot, "trans": pred_trans} - if cfg.TEST.USE_PNP or cfg.TEST.SAVE_RESULTS_ONLY: + if cfg.TEST.USE_PNP or cfg.TEST.SAVE_RESULTS_ONLY or cfg.TEST.USE_DEPTH_REFINE: # TODO: move the pnp/ransac inside forward out_dict.update({"mask": mask, "coor_x": coor_x, "coor_y": coor_y, "coor_z": coor_z, "region": region}) else: diff --git a/core/gdrn_modeling/models/GDRN_Dstream_double_mask.py b/core/gdrn_modeling/models/GDRN_Dstream_double_mask.py index 9d76dc9a6a48f468801e2870016838f26b33e833..cfa13824ee589b4fcefec287cc35ea3031316a48 100644 --- a/core/gdrn_modeling/models/GDRN_Dstream_double_mask.py +++ b/core/gdrn_modeling/models/GDRN_Dstream_double_mask.py @@ -221,7 +221,7 @@ class GDRN_Dstream_DoubleMask(nn.Module): if not do_loss: # test out_dict = {"rot": pred_ego_rot, "trans": pred_trans} - if cfg.TEST.USE_PNP or cfg.TEST.SAVE_RESULTS_ONLY: + if cfg.TEST.USE_PNP or cfg.TEST.SAVE_RESULTS_ONLY or cfg.TEST.USE_DEPTH_REFINE: # TODO: move the pnp/ransac inside forward out_dict.update( { diff --git a/core/gdrn_modeling/models/GDRN_cls.py b/core/gdrn_modeling/models/GDRN_cls.py index 66dc9b25568c11a2d44769886d62d16e0ab29b12..4596fd56534e1c6216ef4bf09ed15f997763cf54 100644 --- a/core/gdrn_modeling/models/GDRN_cls.py +++ b/core/gdrn_modeling/models/GDRN_cls.py @@ -211,7 +211,7 @@ class GDRN_CLS(nn.Module): if not do_loss: # test out_dict = {"rot": pred_ego_rot, "trans": pred_trans} - if cfg.TEST.USE_PNP or cfg.TEST.SAVE_RESULTS_ONLY: + if cfg.TEST.USE_PNP or cfg.TEST.SAVE_RESULTS_ONLY or cfg.TEST.USE_DEPTH_REFINE: # TODO: move the pnp/ransac inside forward out_dict.update({"mask": mask, "coor_x": coor_x, "coor_y": coor_y, "coor_z": coor_z, "region": region}) else: diff --git a/core/gdrn_modeling/models/GDRN_cls2reg.py b/core/gdrn_modeling/models/GDRN_cls2reg.py index 20a566e8dbe161eb4920fd8c2b0da9ef3798c8eb..13bb81414c46c9279942149c3fd7598bb1f86769 100644 --- a/core/gdrn_modeling/models/GDRN_cls2reg.py +++ b/core/gdrn_modeling/models/GDRN_cls2reg.py @@ -211,7 +211,7 @@ class GDRN_CLS2REG(nn.Module): if not do_loss: # test out_dict = {"rot": pred_ego_rot, "trans": pred_trans} - if cfg.TEST.USE_PNP or cfg.TEST.SAVE_RESULTS_ONLY: + if cfg.TEST.USE_PNP or cfg.TEST.SAVE_RESULTS_ONLY or cfg.TEST.USE_DEPTH_REFINE: # TODO: move the pnp/ransac inside forward # TODO: use cls or reg in cfg out_dict.update({"mask": mask, "coor_x": coor_x, "coor_y": coor_y, "coor_z": coor_z, "region": region}) diff --git a/core/gdrn_modeling/models/GDRN_double_mask.py b/core/gdrn_modeling/models/GDRN_double_mask.py index dad1230ce1ea3fd32a9ef613321b7bdcd92ffd08..7973d93a584936c0af9763075e4aee63b9438926 100644 --- a/core/gdrn_modeling/models/GDRN_double_mask.py +++ b/core/gdrn_modeling/models/GDRN_double_mask.py @@ -200,7 +200,7 @@ class GDRN_DoubleMask(nn.Module): if not do_loss: # test out_dict = {"rot": pred_ego_rot, "trans": pred_trans} - if cfg.TEST.USE_PNP or cfg.TEST.SAVE_RESULTS_ONLY: + if cfg.TEST.USE_PNP or cfg.TEST.SAVE_RESULTS_ONLY or cfg.TEST.USE_DEPTH_REFINE: # TODO: move the pnp/ransac inside forward out_dict.update( { diff --git a/core/gdrn_modeling/models/GDRN_no_region.py b/core/gdrn_modeling/models/GDRN_no_region.py index 87c73a81336dacfb847285d481d0098f110265f1..3758999261ccc8440f91c517da9f090bc20ea76e 100644 --- a/core/gdrn_modeling/models/GDRN_no_region.py +++ b/core/gdrn_modeling/models/GDRN_no_region.py @@ -197,7 +197,7 @@ class GDRN_NoRegion(nn.Module): if not do_loss: # test out_dict = {"rot": pred_ego_rot, "trans": pred_trans} - if cfg.TEST.USE_PNP or cfg.TEST.SAVE_RESULTS_ONLY: + if cfg.TEST.USE_PNP or cfg.TEST.SAVE_RESULTS_ONLY or cfg.TEST.USE_DEPTH_REFINE: # TODO: move the pnp/ransac inside forward out_dict.update({"mask": mask, "coor_x": coor_x, "coor_y": coor_y, "coor_z": coor_z, "region": region}) else: