add data process for dpn

configs/catre/NOCS_REAL/aug05_kpsMS_r9d_catreDisR_shared_tspcl_convPerRot_scaleexp_120e_initspd.py

@@ -0,0 +1,135 @@
+_base_ = ["../../_base_/catre_base.py"]
+
+# fix mug mean shape
+OUTPUT_DIR = "output/catre/NOCS_REAL/aug05_kpsMS_r9d_catreDisR_shared_tspcl_convPerRot_scaleexp_120e_initspd"
+INPUT = dict(
+    COLOR_AUG_PROB=0.0,
+    DEPTH_SAMPLE_BALL_RATIO=0.6,
+    # NOTE: used points!
+    BBOX_TYPE_TEST="est",  # from_pose | est | gt | gt_aug (TODO)
+    INIT_POSE_TYPE_TRAIN=["gt_noise"],  # gt_noise | random | canonical
+    NOISE_ROT_STD_TRAIN=(10, 5, 2.5, 1.25),  # randomly choose one
+    NOISE_TRANS_STD_TRAIN=[
+        (0.02, 0.02, 0.02),
+        (0.01, 0.01, 0.01),
+        (0.005, 0.005, 0.005),
+    ],
+    NOISE_SCALE_STD_TRAIN=[
+        (0.01, 0.01, 0.01),
+        (0.005, 0.005, 0.005),
+        (0.002, 0.002, 0.002),
+    ],
+    INIT_POSE_TYPE_TEST="est",  # gt_noise | est | canonical
+    KPS_TYPE="mean_shape",  # bbox_from_scale | mean_shape |fps (abla)
+    WITH_DEPTH=True,
+    AUG_DEPTH=True,
+    WITH_PCL=True,
+    WITH_IMG=False,
+    BP_DEPTH=False,
+    NUM_KPS=1024,
+    NUM_PCL=1024,
+    # augmentation when training
+    BBOX3D_AUG_PROB=0.5,
+    RT_AUG_PROB=0.5,
+    ZERO_CENTER_INPUT=True,
+)
+
+DATALOADER = dict(
+    NUM_WORKERS=24,
+)
+
+SOLVER = dict(
+    IMS_PER_BATCH=32,
+    TOTAL_EPOCHS=120,
+    LR_SCHEDULER_NAME="flat_and_anneal",
+    ANNEAL_METHOD="cosine",  # "cosine"
+    ANNEAL_POINT=0.72,
+    # REL_STEPS=(0.3125, 0.625, 0.9375),
+    OPTIMIZER_CFG=dict(_delete_=True, type="Ranger", lr=1e-4, weight_decay=0),
+    WEIGHT_DECAY=0.0,
+    WARMUP_FACTOR=0.001,
+    WARMUP_ITERS=1000,
+)
+
+DATASETS = dict(
+    TRAIN=("nocs_train_real",),
+    TEST=("nocs_test_real",),
+    INIT_POSE_FILES_TEST=("datasets/NOCS/test_init_poses/init_pose_dualposenet_nocs_real.json",),
+)
+
+MODEL = dict(
+    LOAD_POSES_TEST=True,
+    PIXEL_MEAN=[0.0, 0.0, 0.0],
+    PIXEL_STD=[255.0, 255.0, 255.0],
+    REFINE_SCLAE=True,
+    CATRE=dict(
+        NAME="CATRE_disR_shared",  # used module file name (define different model types)
+        TASK="refine",  # refine | init | init+refine
+        NUM_CLASSES=6,  # only valid for class aware
+        N_ITER_TRAIN=4,
+        N_ITER_TRAIN_WARM_EPOCH=4,  # linearly increase the refine iter from 1 to N_ITER_TRAIN until this epoch
+        N_ITER_TEST=4,
+        PCLNET=dict(
+            FREEZE=False,
+            INIT_CFG=dict(
+                type="point_net",
+                num_points=1024,
+                global_feat=False,
+                feature_transform=True,
+                out_dim=1024,
+            ),
+        ),
+        ## pose head for delta R/T/s
+        ROT_HEAD=dict(
+            ROT_TYPE="ego_rot6d",  # {ego|allo}_rot6d
+            INIT_CFG=dict(
+                type="ConvOutPerRotHead",
+                in_dim=1088,
+                num_layers=2,
+                kernel_size=1,
+                feat_dim=256,
+                norm="GN",  # BN | GN | none
+                num_gn_groups=32,
+                act="gelu",  # relu | lrelu | silu (swish) | gelu | mish
+                num_points=1024 + 1024,
+                rot_dim=3,  # ego_rot6d
+                norm_input=False,
+            ),
+            SCLAE_TYPE="iter_add",
+        ),
+        TS_HEAD=dict(
+            WITH_KPS_FEATURE=False,
+            WITH_INIT_SCALE=True,
+            INIT_CFG=dict(
+                type="FC_TransSizeHead",
+                in_dim=1088 + 3,
+                num_layers=2,
+                feat_dim=256,
+                norm="GN",  # BN | GN | none
+                num_gn_groups=32,
+                act="gelu",  # relu | lrelu | silu (swish) | gelu | mish
+                norm_input=False,
+            ),
+        ),
+        LOSS_CFG=dict(
+            # point matching loss ----------------
+            PM_LOSS_SYM=True,  # use symmetric PM loss
+            PM_NORM_BY_EXTENT=False,  # 1. / extent.max(1, keepdim=True)[0]
+            # if False, the trans loss is in point matching loss
+            PM_R_ONLY=True,  # only do R loss in PM
+            PM_WITH_SCALE=True,
+            PM_LW=1.0,
+            # rot loss --------------
+            ROT_LOSS_TYPE="angular",  # angular | L2
+            ROT_LW=1.0,
+            ROT_YAXIS_LOSS_TYPE="L1",
+            # trans loss -----------
+            TRANS_LOSS_TYPE="L1",
+            TRANS_LOSS_DISENTANGLE=True,
+            TRANS_LW=1.0,
+            # scale loss ----------------------------------
+            SCALE_LOSS_TYPE="L1",
+            SCALE_LW=1.0,
+        ),
+    ),
+)

core/catre/tools/prepare_dualposenet_init_results.py

@@ -0,0 +1,106 @@
+import numpy as np
+import mmcv
+import os.path as osp
+import sys
+from tqdm import tqdm
+import setproctitle
+
+cur_dir = osp.dirname(osp.abspath(__file__))
+PROJ_ROOT = osp.normpath(osp.join(cur_dir, "../../../"))
+sys.path.insert(0, PROJ_ROOT)
+from lib.pysixd import inout, misc
+from lib.utils.mask_utils import binary_mask_to_rle
+
+setproctitle.setproctitle(osp.basename(__file__).split(".")[0])
+
+data_root = osp.join(PROJ_ROOT, "datasets/NOCS")
+
+# original dualposenet results
+dpn_pose_dir = osp.join(PROJ_ROOT, "datasets/NOCS/dualposenet_results")  # # provided by the authors
+dpn_seg_dir = osp.join(PROJ_ROOT, "datasets/NOCS/dualposenet_results/segmentation_results/REAL275")
+
+# our format
+init_pose_dir = osp.join(data_root, "test_init_poses")
+mmcv.mkdir_or_exist(init_pose_dir)
+init_pose_path = osp.join(init_pose_dir, "init_pose_dualposenet_nocs_real.json")
+
+
+if __name__ == "__main__":
+    results = {}
+
+    CACHED = False
+    if CACHED:
+        results = mmcv.load(init_pose_path)
+    else:
+        dpn_pose_path = osp.join(dpn_pose_dir, "REAL275_results_mvmug.pkl")
+        dpn_pose_res = mmcv.load(dpn_pose_path)
+
+        num_total = 0
+        for idx, preds in enumerate(tqdm(dpn_pose_res)):
+            bboxes = preds["pred_bboxes"]
+            scores = preds["pred_scores"]
+            poses = preds["pred_RTs"][:, :3]
+            pred_scales = preds["pred_scales"]
+            class_ids = preds["pred_class_ids"]
+            mug_handles = preds["gt_handle_visibility"]
+
+            scene_id, im_id = preds["image_path"].split("/")[-2:]
+            scene_im_id = f"{scene_id}/{im_id}"
+
+            seg_path = osp.join(dpn_seg_dir, f"results_test_{scene_id}_{im_id}.pkl")
+            assert osp.exists(seg_path)
+
+            masks = mmcv.load(seg_path)["pred_masks"].astype("int")  # bool -> int
+            assert masks.shape[2] == len(class_ids)
+
+            results[scene_im_id] = []
+            i = 0
+            for class_id, pose, scale, score, bbox, mug_handle in zip(
+                class_ids, poses, pred_scales, scores, bboxes, mug_handles
+            ):
+                # [sR -> R], normed_scale -> scale
+                R = pose[:3, :3]
+                nocs_scale = pow(np.linalg.det(R), 1 / 3)
+                abs_scale = scale * nocs_scale
+                pose[:3, :3] = R / nocs_scale
+                # mask2rle
+                mask = masks[:, :, i]
+                mask_rle = binary_mask_to_rle(mask)
+                y1, x1, y2, x2 = bbox.tolist()
+                bbox = [x1, y1, x2, y2]
+                cur_res = {
+                    "obj_id": int(class_id),
+                    "pose_est": pose.tolist(),
+                    "scale_est": abs_scale.tolist(),
+                    "bbox_est": bbox,
+                    "score": float(score),
+                    "mug_handle": int(mug_handle),
+                    "segmentation": mask_rle,
+                }
+                results[scene_im_id].append(cur_res)
+                i += 1
+
+        print(init_pose_path)
+        # inout.save_json(init_pose_path, results, sort=False)
+
+    VIS = False
+    if VIS:
+        from core.utils.data_utils import read_image_mmcv
+        from lib.utils.mask_utils import cocosegm2mask
+        from lib.vis_utils.image import grid_show, heatmap
+        from core.catre.engine.test_utils import get_3d_bbox
+        import ref
+
+        for scene_im_id, r in results.items():
+            img_path = f"datasets/NOCS/REAL/real_test/{scene_im_id}_color.png"
+            img = read_image_mmcv(img_path, format="BGR")
+            K = ref.nocs.real_intrinsics
+            anno = r[0]
+            imH, imW = img.shape[:2]
+            mask = cocosegm2mask(anno["segmentation"], imH, imW)
+            pose = np.array(anno["pose_est"]).reshape(3, 4)
+            scale = np.array(anno["scale_est"])
+            bbox = get_3d_bbox(scale).transpose()
+            kpts_2d = misc.project_pts(bbox, K, pose[:, :3], pose[:, 3])
+            img_vis_kpts2d = misc.draw_projected_box3d(img.copy(), kpts_2d)
+            grid_show([img[:, :, ::-1], img_vis_kpts2d, mask], row=3, col=1)