From 93ec7d4a68b4bafa8765eece26b651d727d42216 Mon Sep 17 00:00:00 2001 From: asamsone <anna.samsonenko@etu.ec-lyon.fr> Date: Fri, 23 Feb 2024 12:49:21 +0100 Subject: [PATCH] added changes for fruitbin --- ...DMask_amodalClipBox_classAware_fruitbin.py | 29 ++- .../banana1.py | 22 ++- .../lemon2.py | 6 +- .../pear2.py | 6 +- .../{002_master_chef_can.py => banana1.py} | 22 ++- .../lemon2.py | 6 +- .../pear2.py | 6 +- ...0_epochs_fruitbin_pbr_fruitbin_bop_test.py | 8 +- core/gdrn_modeling/datasets/data_loader.py | 6 +- .../datasets/data_loader_online.py | 4 +- .../gdrn_modeling/datasets/dataset_factory.py | 17 ++ .../datasets/fruitbin_bop_test.py | 94 +++++----- core/gdrn_modeling/datasets/fruitbin_d2.py | 165 +++++++++--------- core/gdrn_modeling/datasets/fruitbin_pbr.py | 85 +++++---- core/gdrn_modeling/main_gdrn.py | 1 - .../fruitbin/convert_det_to_our_format.py | 69 -------- .../tools/fruitbin/fruitbin_1_compute_fps.py | 45 +++++ det/yolox/data/datasets/dataset_factory.py | 6 + det/yolox/data/datasets/fruitbin_bop_test.py | 94 +++++----- det/yolox/data/datasets/fruitbin_d2.py | 163 ++++++++--------- det/yolox/data/datasets/fruitbin_pbr.py | 89 +++++----- ref/__init__.py | 2 +- ref/fruitbin.py | 67 +++---- 23 files changed, 495 insertions(+), 517 deletions(-) rename configs/gdrn/fruitbinSO/convnext_AugCosyAAEGray_DMask_amodalClipBox_fruitbin/{002_master_chef_can.py => banana1.py} (87%) delete mode 100644 core/gdrn_modeling/tools/fruitbin/convert_det_to_our_format.py create mode 100644 core/gdrn_modeling/tools/fruitbin/fruitbin_1_compute_fps.py diff --git a/configs/gdrn/fruitbin/convnext_a6_AugCosyAAEGray_BG05_mlL1_DMask_amodalClipBox_classAware_fruitbin.py b/configs/gdrn/fruitbin/convnext_a6_AugCosyAAEGray_BG05_mlL1_DMask_amodalClipBox_classAware_fruitbin.py index b167a6c..339caa1 100644 --- a/configs/gdrn/fruitbin/convnext_a6_AugCosyAAEGray_BG05_mlL1_DMask_amodalClipBox_classAware_fruitbin.py +++ b/configs/gdrn/fruitbin/convnext_a6_AugCosyAAEGray_BG05_mlL1_DMask_amodalClipBox_classAware_fruitbin.py @@ -1,7 +1,7 @@ # about 3 days _base_ = ["../../_base_/gdrn_base.py"] -OUTPUT_DIR = "output/gdrn/ycbv/convnext_a6_AugCosyAAEGray_BG05_mlL1_DMask_amodalClipBox_classAware_ycbv" +OUTPUT_DIR = "output/gdrn/fruitbin/convnext_a6_AugCosyAAEGray_BG05_mlL1_DMask_amodalClipBox_classAware_fruitbin" INPUT = dict( DZI_PAD_SCALE=1.5, TRUNCATE_FG=True, @@ -32,8 +32,8 @@ INPUT = dict( ) SOLVER = dict( - IMS_PER_BATCH=48, - TOTAL_EPOCHS=40, # 10 + IMS_PER_BATCH=36, + TOTAL_EPOCHS=1, # 10 LR_SCHEDULER_NAME="flat_and_anneal", ANNEAL_METHOD="cosine", # "cosine" ANNEAL_POINT=0.72, @@ -44,15 +44,12 @@ SOLVER = dict( ) DATASETS = dict( - TRAIN=("ycbv_train_real", "ycbv_train_pbr"), - TEST=("ycbv_test",), - DET_FILES_TEST=("datasets/BOP_DATASETS/ycbv/test/test_bboxes/yolox_x_640_ycbv_real_pbr_ycbv_bop_test.json",), + TRAIN=("fruitbin_train_pbr",), + TEST=("fruitbin_test",), + DET_FILES_TEST=("datasets/BOP_DATASETS/fruitbin/test/test_bboxes/yolox_x_640_fruitbin_real_pbr_fruitbin_bop_test.json",), SYM_OBJS=[ - "024_bowl", - "036_wood_block", - "051_large_clamp", - "052_extra_large_clamp", - "061_foam_brick", + "lemon2", + "orange2", ], # used for custom evalutor ) @@ -66,11 +63,11 @@ MODEL = dict( LOAD_DETS_TEST=True, PIXEL_MEAN=[0.0, 0.0, 0.0], PIXEL_STD=[255.0, 255.0, 255.0], - BBOX_TYPE="AMODAL_CLIP", # VISIB or AMODAL + BBOX_TYPE="VISIB", # VISIB or AMODAL POSE_NET=dict( NAME="GDRN_double_mask", XYZ_ONLINE=True, - NUM_CLASSES=21, + NUM_CLASSES=8, BACKBONE=dict( FREEZE=False, PRETRAINED="timm", @@ -112,8 +109,8 @@ MODEL = dict( MASK_LW=1.0, # full mask loss --------------------------- FULL_MASK_LOSS_TYPE="L1", # L1 | BCE | CE - FULL_MASK_LW=1.0, - # region loss ------------------------- + FULL_MASK_LW=0.0, + REGION_LOSS_TYPE="CE", # CE REGION_LOSS_MASK_GT="visib", # trunc | visib | obj REGION_LW=1.0, @@ -132,7 +129,7 @@ MODEL = dict( ) VAL = dict( - DATASET_NAME="ycbv", + DATASET_NAME="fruitbin", SPLIT_TYPE="", SCRIPT_PATH="lib/pysixd/scripts/eval_pose_results_more.py", TARGETS_FILENAME="test_targets_bop19.json", diff --git a/configs/gdrn/fruitbinPbrSOr/convnext_AugCosyAAEGray_DMask_amodalClipBox_fruitbin/banana1.py b/configs/gdrn/fruitbinPbrSOr/convnext_AugCosyAAEGray_DMask_amodalClipBox_fruitbin/banana1.py index a7f363f..56ff018 100644 --- a/configs/gdrn/fruitbinPbrSOr/convnext_AugCosyAAEGray_DMask_amodalClipBox_fruitbin/banana1.py +++ b/configs/gdrn/fruitbinPbrSOr/convnext_AugCosyAAEGray_DMask_amodalClipBox_fruitbin/banana1.py @@ -1,6 +1,6 @@ _base_ = ["../../../_base_/gdrn_base.py"] -OUTPUT_DIR = "output/gdrn/ycbvPbrSO/convnext_AugCosyAAEGray_DMask_amodalClipBox_ycbv/002_master_chef_can" +OUTPUT_DIR = "output/gdrn/fruitbinPbrSO/convnext_AugCosyAAEGray_DMask_amodalClipBox_fruitbin/banana1" INPUT = dict( DZI_PAD_SCALE=1.5, TRUNCATE_FG=True, @@ -42,16 +42,14 @@ SOLVER = dict( WARMUP_ITERS=1000, ) + DATASETS = dict( - TRAIN=("ycbv_002_master_chef_can_train_pbr",), - TEST=("ycbv_test",), - DET_FILES_TEST=("datasets/BOP_DATASETS/ycbv/test/test_bboxes/yolox_x_640_ycbv_pbr_ycbv_bop_test.json",), + TRAIN=("fruitbin_train_pbr"), + TEST=("fruitbin",), + # DET_FILES_TEST=("datasets/BOP_DATASETS/fruitbin/test/test_bboxes/yolox_x_640_ycbv_pbr_ycbv_bop_test.json",), SYM_OBJS=[ - "024_bowl", - "036_wood_block", - "051_large_clamp", - "052_extra_large_clamp", - "061_foam_brick", + "lemon2", + "orange2" ], # used for custom evalutor ) @@ -65,7 +63,7 @@ MODEL = dict( LOAD_DETS_TEST=True, PIXEL_MEAN=[0.0, 0.0, 0.0], PIXEL_STD=[255.0, 255.0, 255.0], - BBOX_TYPE="AMODAL_CLIP", # VISIB or AMODAL + BBOX_TYPE="VISIB", # VISIB or AMODAL POSE_NET=dict( NAME="GDRN_double_mask", XYZ_ONLINE=True, @@ -108,7 +106,7 @@ MODEL = dict( # full mask loss --------------------------- FULL_MASK_LOSS_TYPE="L1", # L1 | BCE | CE FULL_MASK_LW=1.0, - # region loss ------------------------- + REGION_LOSS_TYPE="CE", # CE REGION_LOSS_MASK_GT="visib", # trunc | visib | obj REGION_LW=1.0, @@ -127,7 +125,7 @@ MODEL = dict( ) VAL = dict( - DATASET_NAME="ycbv", + DATASET_NAME="fruitbin", SPLIT_TYPE="", SCRIPT_PATH="lib/pysixd/scripts/eval_pose_results_more.py", TARGETS_FILENAME="test_targets_bop19.json", diff --git a/configs/gdrn/fruitbinPbrSOr/convnext_AugCosyAAEGray_DMask_amodalClipBox_fruitbin/lemon2.py b/configs/gdrn/fruitbinPbrSOr/convnext_AugCosyAAEGray_DMask_amodalClipBox_fruitbin/lemon2.py index 59dc1d6..068ba44 100644 --- a/configs/gdrn/fruitbinPbrSOr/convnext_AugCosyAAEGray_DMask_amodalClipBox_fruitbin/lemon2.py +++ b/configs/gdrn/fruitbinPbrSOr/convnext_AugCosyAAEGray_DMask_amodalClipBox_fruitbin/lemon2.py @@ -1,3 +1,3 @@ -_base_ = "./002_master_chef_can.py" -OUTPUT_DIR = "output/gdrn/ycbvPbrSO/convnext_AugCosyAAEGray_DMask_amodalClipBox_ycbv/003_cracker_box" -DATASETS = dict(TRAIN=("ycbv_003_cracker_box_train_pbr",)) +_base_ = "./banana1.py" +OUTPUT_DIR = "output/gdrn/fruitbinPbrSO/convnext_AugCosyAAEGray_DMask_amodalClipBox_fruitbin/lemon2" +DATASETS = dict(TRAIN=("fruitbin_lemon2",)) diff --git a/configs/gdrn/fruitbinPbrSOr/convnext_AugCosyAAEGray_DMask_amodalClipBox_fruitbin/pear2.py b/configs/gdrn/fruitbinPbrSOr/convnext_AugCosyAAEGray_DMask_amodalClipBox_fruitbin/pear2.py index d03bd4d..a905a8e 100644 --- a/configs/gdrn/fruitbinPbrSOr/convnext_AugCosyAAEGray_DMask_amodalClipBox_fruitbin/pear2.py +++ b/configs/gdrn/fruitbinPbrSOr/convnext_AugCosyAAEGray_DMask_amodalClipBox_fruitbin/pear2.py @@ -1,3 +1,3 @@ -_base_ = "./002_master_chef_can.py" -OUTPUT_DIR = "output/gdrn/ycbvPbrSO/convnext_AugCosyAAEGray_DMask_amodalClipBox_ycbv/004_sugar_box" -DATASETS = dict(TRAIN=("ycbv_004_sugar_box_train_pbr",)) +_base_ = "./banana1.py" +OUTPUT_DIR = "output/gdrn/fruitbinPbrSO/convnext_AugCosyAAEGray_DMask_amodalClipBox_fruitbin/pear2" +DATASETS = dict(TRAIN=("fruitbin_pear2",)) diff --git a/configs/gdrn/fruitbinSO/convnext_AugCosyAAEGray_DMask_amodalClipBox_fruitbin/002_master_chef_can.py b/configs/gdrn/fruitbinSO/convnext_AugCosyAAEGray_DMask_amodalClipBox_fruitbin/banana1.py similarity index 87% rename from configs/gdrn/fruitbinSO/convnext_AugCosyAAEGray_DMask_amodalClipBox_fruitbin/002_master_chef_can.py rename to configs/gdrn/fruitbinSO/convnext_AugCosyAAEGray_DMask_amodalClipBox_fruitbin/banana1.py index 8551bf5..dbfda48 100644 --- a/configs/gdrn/fruitbinSO/convnext_AugCosyAAEGray_DMask_amodalClipBox_fruitbin/002_master_chef_can.py +++ b/configs/gdrn/fruitbinSO/convnext_AugCosyAAEGray_DMask_amodalClipBox_fruitbin/banana1.py @@ -1,6 +1,6 @@ _base_ = ["../../../_base_/gdrn_base.py"] -OUTPUT_DIR = "output/gdrn/ycbvSO/convnext_AugCosyAAEGray_DMask_amodalClipBox_ycbv/002_master_chef_can" +OUTPUT_DIR = "output/gdrn/fruitbinPbrSO/convnext_AugCosyAAEGray_DMask_amodalClipBox_fruitbin/banana1" INPUT = dict( DZI_PAD_SCALE=1.5, TRUNCATE_FG=True, @@ -42,16 +42,14 @@ SOLVER = dict( WARMUP_ITERS=1000, ) + DATASETS = dict( - TRAIN=("ycbv_002_master_chef_can_train_pbr", "ycbv_002_master_chef_can_train_real"), - TEST=("ycbv_test",), - DET_FILES_TEST=("datasets/BOP_DATASETS/ycbv/test/test_bboxes/yolox_x_640_ycbv_real_pbr_ycbv_bop_test.json",), + TRAIN=("fruitbin_train_pbr"), + TEST=("fruitbin",), + DET_FILES_TEST=("datasets/BOP_DATASETS/fruitbin/test/test_bboxes/yolox_x_640_ycbv_pbr_ycbv_bop_test.json",), SYM_OBJS=[ - "024_bowl", - "036_wood_block", - "051_large_clamp", - "052_extra_large_clamp", - "061_foam_brick", + "lemon2", + "orange2" ], # used for custom evalutor ) @@ -65,7 +63,7 @@ MODEL = dict( LOAD_DETS_TEST=True, PIXEL_MEAN=[0.0, 0.0, 0.0], PIXEL_STD=[255.0, 255.0, 255.0], - BBOX_TYPE="AMODAL_CLIP", # VISIB or AMODAL + BBOX_TYPE="VISIB", # VISIB or AMODAL POSE_NET=dict( NAME="GDRN_double_mask", XYZ_ONLINE=True, @@ -108,7 +106,7 @@ MODEL = dict( # full mask loss --------------------------- FULL_MASK_LOSS_TYPE="L1", # L1 | BCE | CE FULL_MASK_LW=1.0, - # region loss ------------------------- + REGION_LOSS_TYPE="CE", # CE REGION_LOSS_MASK_GT="visib", # trunc | visib | obj REGION_LW=1.0, @@ -127,7 +125,7 @@ MODEL = dict( ) VAL = dict( - DATASET_NAME="ycbv", + DATASET_NAME="fruitbin", SPLIT_TYPE="", SCRIPT_PATH="lib/pysixd/scripts/eval_pose_results_more.py", TARGETS_FILENAME="test_targets_bop19.json", diff --git a/configs/gdrn/fruitbinSO/convnext_AugCosyAAEGray_DMask_amodalClipBox_fruitbin/lemon2.py b/configs/gdrn/fruitbinSO/convnext_AugCosyAAEGray_DMask_amodalClipBox_fruitbin/lemon2.py index 23e99a3..068ba44 100644 --- a/configs/gdrn/fruitbinSO/convnext_AugCosyAAEGray_DMask_amodalClipBox_fruitbin/lemon2.py +++ b/configs/gdrn/fruitbinSO/convnext_AugCosyAAEGray_DMask_amodalClipBox_fruitbin/lemon2.py @@ -1,3 +1,3 @@ -_base_ = "./002_master_chef_can.py" -OUTPUT_DIR = "output/gdrn/ycbvSO/convnext_AugCosyAAEGray_DMask_amodalClipBox_ycbv/003_cracker_box" -DATASETS = dict(TRAIN=("ycbv_003_cracker_box_train_pbr", "ycbv_003_cracker_box_train_real")) +_base_ = "./banana1.py" +OUTPUT_DIR = "output/gdrn/fruitbinPbrSO/convnext_AugCosyAAEGray_DMask_amodalClipBox_fruitbin/lemon2" +DATASETS = dict(TRAIN=("fruitbin_lemon2",)) diff --git a/configs/gdrn/fruitbinSO/convnext_AugCosyAAEGray_DMask_amodalClipBox_fruitbin/pear2.py b/configs/gdrn/fruitbinSO/convnext_AugCosyAAEGray_DMask_amodalClipBox_fruitbin/pear2.py index efa07cb..a905a8e 100644 --- a/configs/gdrn/fruitbinSO/convnext_AugCosyAAEGray_DMask_amodalClipBox_fruitbin/pear2.py +++ b/configs/gdrn/fruitbinSO/convnext_AugCosyAAEGray_DMask_amodalClipBox_fruitbin/pear2.py @@ -1,3 +1,3 @@ -_base_ = "./002_master_chef_can.py" -OUTPUT_DIR = "output/gdrn/ycbvSO/convnext_AugCosyAAEGray_DMask_amodalClipBox_ycbv/004_sugar_box" -DATASETS = dict(TRAIN=("ycbv_004_sugar_box_train_pbr", "ycbv_004_sugar_box_train_real")) +_base_ = "./banana1.py" +OUTPUT_DIR = "output/gdrn/fruitbinPbrSO/convnext_AugCosyAAEGray_DMask_amodalClipBox_fruitbin/pear2" +DATASETS = dict(TRAIN=("fruitbin_pear2",)) diff --git a/configs/yolox/bop_pbr/yolox_x_640_augCozyAAEhsv_ranger_30_epochs_fruitbin_pbr_fruitbin_bop_test.py b/configs/yolox/bop_pbr/yolox_x_640_augCozyAAEhsv_ranger_30_epochs_fruitbin_pbr_fruitbin_bop_test.py index 729c297..0b49f49 100644 --- a/configs/yolox/bop_pbr/yolox_x_640_augCozyAAEhsv_ranger_30_epochs_fruitbin_pbr_fruitbin_bop_test.py +++ b/configs/yolox/bop_pbr/yolox_x_640_augCozyAAEhsv_ranger_30_epochs_fruitbin_pbr_fruitbin_bop_test.py @@ -20,16 +20,16 @@ train.amp.enabled = True model.backbone.depth = 1.33 model.backbone.width = 1.25 -model.head.num_classes = 21 +model.head.num_classes = 8 train.init_checkpoint = "pretrained_models/yolox/yolox_x.pth" # datasets -DATASETS.TRAIN = ["ycbv_train_pbr"] -DATASETS.TEST = ["ycbv_bop_test"] +DATASETS.TRAIN = ["fruitbin_train_pbr"] +DATASETS.TEST = ["fruitbin_bop_test"] dataloader.train.dataset.lst.names = DATASETS.TRAIN -dataloader.train.total_batch_size = 32 +dataloader.train.total_batch_size = 15 # color aug dataloader.train.aug_wrapper.COLOR_AUG_PROB = 0.8 diff --git a/core/gdrn_modeling/datasets/data_loader.py b/core/gdrn_modeling/datasets/data_loader.py index cbc1eb2..7cc5527 100644 --- a/core/gdrn_modeling/datasets/data_loader.py +++ b/core/gdrn_modeling/datasets/data_loader.py @@ -64,10 +64,10 @@ def transform_instance_annotations(annotation, transforms, image_size, *, keypoi """ im_H, im_W = image_size bbox = BoxMode.convert(annotation["bbox"], annotation["bbox_mode"], BoxMode.XYXY_ABS) - bbox_obj = BoxMode.convert(annotation["bbox_obj"], annotation["bbox_mode"], BoxMode.XYXY_ABS) + # bbox_obj = BoxMode.convert(annotation["bbox_obj"], annotation["bbox_mode"], BoxMode.XYXY_ABS) # Note that bbox is 1d (per-instance bounding box) annotation["bbox"] = np.array(transforms.apply_box([bbox])[0]) - annotation["bbox_obj"] = np.array(transforms.apply_box([bbox_obj])[0]) + # annotation["bbox_obj"] = np.array(transforms.apply_box([bbox_obj])[0]) annotation["bbox_mode"] = BoxMode.XYXY_ABS if "segmentation" in annotation: @@ -922,6 +922,8 @@ def build_gdrn_test_loader(cfg, dataset_name, train_objs=None, sampler=None, bat # load test detection results if cfg.MODEL.LOAD_DETS_TEST: det_files = cfg.DATASETS.DET_FILES_TEST + print("len(cfg.DATASETS.TEST) : ", len(cfg.DATASETS.TEST)) + print("len(det_files) : ", len(det_files)) assert len(cfg.DATASETS.TEST) == len(det_files) dataset_dicts = load_detections_into_dataset( dataset_name, diff --git a/core/gdrn_modeling/datasets/data_loader_online.py b/core/gdrn_modeling/datasets/data_loader_online.py index 257985a..37d4d3c 100644 --- a/core/gdrn_modeling/datasets/data_loader_online.py +++ b/core/gdrn_modeling/datasets/data_loader_online.py @@ -58,10 +58,10 @@ def transform_instance_annotations(annotation, transforms, image_size, *, keypoi """ im_H, im_W = image_size bbox = BoxMode.convert(annotation["bbox"], annotation["bbox_mode"], BoxMode.XYXY_ABS) - bbox_obj = BoxMode.convert(annotation["bbox_obj"], annotation["bbox_mode"], BoxMode.XYXY_ABS) + # bbox_obj = BoxMode.convert(annotation["bbox_obj"], annotation["bbox_mode"], BoxMode.XYXY_ABS) # Note that bbox is 1d (per-instance bounding box) annotation["bbox"] = np.array(transforms.apply_box([bbox])[0]) - annotation["bbox_obj"] = np.array(transforms.apply_box([bbox_obj])[0]) + # annotation["bbox_obj"] = np.array(transforms.apply_box([bbox_obj])[0]) annotation["bbox_mode"] = BoxMode.XYXY_ABS if "segmentation" in annotation: diff --git a/core/gdrn_modeling/datasets/dataset_factory.py b/core/gdrn_modeling/datasets/dataset_factory.py index c60929e..9a7b745 100644 --- a/core/gdrn_modeling/datasets/dataset_factory.py +++ b/core/gdrn_modeling/datasets/dataset_factory.py @@ -27,6 +27,10 @@ from core.gdrn_modeling.datasets import ( itodd_pbr, itodd_bop_test, itodd_d2, + fruitbin_d2, + fruitbin_pbr, + fruitbin_bop_test, + ) @@ -58,6 +62,10 @@ _DSET_MOD_NAMES = [ "itodd_pbr", "itodd_bop_test", "itodd_d2", + "fruitbin_d2", + "fruitbin_pbr", + "fruitbin_bop_test", + ] logger = logging.getLogger(__name__) @@ -87,13 +95,20 @@ def register_datasets_in_cfg(cfg): "TRAIN2", "TRAIN_SYN_SUP", ]: + print(cfg.DATASETS.get(split, [])) for name in cfg.DATASETS.get(split, []): if name in DatasetCatalog.list(): continue registered = False # try to find in pre-defined datasets # NOTE: it is better to let all datasets pre-refined + print(_DSET_MOD_NAMES) for _mod_name in _DSET_MOD_NAMES: + print("-") + print("get_available_datasets") + print(get_available_datasets(_mod_name)) + print("name") + print(name) if name in get_available_datasets(_mod_name): register_dataset(_mod_name, name, data_cfg=None) registered = True @@ -104,6 +119,8 @@ def register_datasets_in_cfg(cfg): """load data_cfg and mod_name from file cfg.DATA_CFG[name] = 'path_to_cfg' """ + # print(cfg.DATA_CFG) + # ffs assert "DATA_CFG" in cfg and name in cfg.DATA_CFG, "no cfg.DATA_CFG.{}".format(name) assert osp.exists(cfg.DATA_CFG[name]) data_cfg = mmcv.load(cfg.DATA_CFG[name]) diff --git a/core/gdrn_modeling/datasets/fruitbin_bop_test.py b/core/gdrn_modeling/datasets/fruitbin_bop_test.py index 43dd74b..956111b 100644 --- a/core/gdrn_modeling/datasets/fruitbin_bop_test.py +++ b/core/gdrn_modeling/datasets/fruitbin_bop_test.py @@ -25,8 +25,8 @@ logger = logging.getLogger(__name__) DATASETS_ROOT = osp.normpath(osp.join(PROJ_ROOT, "datasets")) -class YCBV_BOP_TEST_Dataset: - """ycbv bop test.""" +class FRUITBIN_BOP_TEST_Dataset: + """fruitbin bop test.""" def __init__(self, data_cfg): """ @@ -43,8 +43,8 @@ class YCBV_BOP_TEST_Dataset: self.ann_file = data_cfg["ann_file"] # json file with scene_id and im_id items - self.dataset_root = data_cfg["dataset_root"] # BOP_DATASETS/ycbv/test - self.models_root = data_cfg["models_root"] # BOP_DATASETS/ycbv/models + self.dataset_root = data_cfg["dataset_root"] # BOP_DATASETS/fruitbin/test + self.models_root = data_cfg["models_root"] # BOP_DATASETS/fruitbin/models self.scale_to_meter = data_cfg["scale_to_meter"] # 0.001 self.with_masks = data_cfg["with_masks"] # True (load masks but may not use it) @@ -60,7 +60,7 @@ class YCBV_BOP_TEST_Dataset: ################################################## # NOTE: careful! Only the selected objects - self.cat_ids = [cat_id for cat_id, obj_name in ref.ycbv.id2obj.items() if obj_name in self.objs] + self.cat_ids = [cat_id for cat_id, obj_name in ref.fruitbin.id2obj.items() if obj_name in self.objs] # map selected objs to [0, num_objs-1] self.cat2label = {v: i for i, v in enumerate(self.cat_ids)} # id_map self.label2cat = {label: cat for cat, label in self.cat2label.items()} @@ -127,10 +127,10 @@ class YCBV_BOP_TEST_Dataset: for scene_id, im_id in tqdm(scene_im_ids): str_im_id = str(im_id) scene_root = osp.join(self.dataset_root, f"{scene_id:06d}") - rgb_path = osp.join(scene_root, "rgb/{:06d}.png").format(im_id) + rgb_path = osp.join(scene_root, "rgb/{:d}.png").format(im_id) assert osp.exists(rgb_path), rgb_path - depth_path = osp.join(scene_root, "depth/{:06d}.png".format(im_id)) + depth_path = osp.join(scene_root, "depth/{:d}.png".format(im_id)) scene_id = int(rgb_path.split("/")[-3]) @@ -152,7 +152,7 @@ class YCBV_BOP_TEST_Dataset: insts = [] for anno_i, anno in enumerate(gt_dicts[scene_id][str_im_id]): obj_id = anno["obj_id"] - if ref.ycbv.id2obj[obj_id] not in self.select_objs: + if ref.fruitbin.id2obj[obj_id] not in self.select_objs: continue cur_label = self.cat2label[obj_id] # 0-based label R = np.array(anno["cam_R_m2c"], dtype="float32").reshape(3, 3) @@ -164,22 +164,22 @@ class YCBV_BOP_TEST_Dataset: proj = proj[:2] / proj[2] bbox_visib = gt_info_dicts[scene_id][str_im_id][anno_i]["bbox_visib"] - bbox_obj = gt_info_dicts[scene_id][str_im_id][anno_i]["bbox_obj"] + # bbox_obj = gt_info_dicts[scene_id][str_im_id][anno_i]["bbox_obj"] x1, y1, w, h = bbox_visib if self.filter_invalid: if h <= 1 or w <= 1: self.num_instances_without_valid_box += 1 continue - mask_file = osp.join( - scene_root, - "mask/{:06d}_{:06d}.png".format(im_id, anno_i), - ) + # mask_file = osp.join( + # scene_root, + # "mask/{:06d}_{:06d}.png".format(im_id, anno_i), + # ) mask_visib_file = osp.join( scene_root, - "mask_visib/{:06d}_{:06d}.png".format(im_id, anno_i), + "mask_visib/{:d}.png".format(im_id, anno_i), ) - assert osp.exists(mask_file), mask_file + # assert osp.exists(mask_file), mask_file assert osp.exists(mask_visib_file), mask_visib_file # load mask visib mask_single = mmcv.imread(mask_visib_file, "unchanged") @@ -190,21 +190,21 @@ class YCBV_BOP_TEST_Dataset: mask_rle = binary_mask_to_rle(mask_single, compressed=True) # load mask full - mask_full = mmcv.imread(mask_file, "unchanged") - mask_full = mask_full.astype("bool") - mask_full_rle = binary_mask_to_rle(mask_full, compressed=True) + # mask_full = mmcv.imread(mask_file, "unchanged") + # mask_full = mask_full.astype("bool") + # mask_full_rle = binary_mask_to_rle(mask_full, compressed=True) inst = { "category_id": cur_label, # 0-based label "bbox": bbox_visib, - "bbox_obj": bbox_obj, + # "bbox_obj": bbox_obj, "bbox_mode": BoxMode.XYWH_ABS, "pose": pose, "quat": quat, "trans": t, "centroid_2d": proj, # absolute (cx, cy) "segmentation": mask_rle, - "mask_full": mask_full_rle, # TODO: load as mask_full, rle + # "mask_full": mask_full_rle, # TODO: load as mask_full, rle } model_info = self.models_info[str(obj_id)] @@ -261,7 +261,7 @@ class YCBV_BOP_TEST_Dataset: model = inout.load_ply( osp.join( self.models_root, - f"obj_{ref.ycbv.obj2id[obj_name]:06d}.ply", + f"obj_{ref.fruitbin.obj2id[obj_name]:06d}.ply", ), vertex_scale=self.scale_to_meter, ) @@ -281,7 +281,7 @@ class YCBV_BOP_TEST_Dataset: ########### register datasets ############################################################ -def get_ycbv_metadata(obj_names, ref_key): +def get_fruitbin_metadata(obj_names, ref_key): """task specific metadata.""" data_ref = ref.__dict__[ref_key] @@ -304,49 +304,49 @@ def get_ycbv_metadata(obj_names, ref_key): ################################################################################ -SPLITS_YCBV = dict( - ycbv_bop_test=dict( - name="ycbv_bop_test", - dataset_root=osp.join(DATASETS_ROOT, "BOP_DATASETS/ycbv/test"), - models_root=osp.join(DATASETS_ROOT, "BOP_DATASETS/ycbv/models"), - objs=ref.ycbv.objects, # selected objects - ann_file=osp.join(DATASETS_ROOT, "BOP_DATASETS/ycbv/test_targets_bop19.json"), +SPLITS_FRUITBIN = dict( + fruitbin_bop_test=dict( + name="fruitbin_bop_test", + dataset_root=osp.join(DATASETS_ROOT, "BOP_DATASETS/fruitbin/test"), + models_root=osp.join(DATASETS_ROOT, "BOP_DATASETS/fruitbin/models"), + objs=ref.fruitbin.objects, # selected objects + ann_file=osp.join(DATASETS_ROOT, "BOP_DATASETS/fruitbin/test_targets_bop19.json"), scale_to_meter=0.001, with_masks=True, # (load masks but may not use it) with_depth=True, # (load depth path here, but may not use it) height=480, width=640, cache_dir=osp.join(PROJ_ROOT, ".cache"), - use_cache=True, + use_cache=False, num_to_load=-1, filter_invalid=False, - ref_key="ycbv", + ref_key="fruitbin", ) ) # single objs (num_class is from all objs) -for obj in ref.ycbv.objects: - name = "ycbv_bop_{}_test".format(obj) +for obj in ref.fruitbin.objects: + name = "fruitbin_bop_{}_test".format(obj) select_objs = [obj] - if name not in SPLITS_YCBV: - SPLITS_YCBV[name] = dict( + if name not in SPLITS_FRUITBIN: + SPLITS_FRUITBIN[name] = dict( name=name, - dataset_root=osp.join(DATASETS_ROOT, "BOP_DATASETS/ycbv/test"), - models_root=osp.join(DATASETS_ROOT, "BOP_DATASETS/ycbv/models"), + dataset_root=osp.join(DATASETS_ROOT, "BOP_DATASETS/fruitbin/test"), + models_root=osp.join(DATASETS_ROOT, "BOP_DATASETS/fruitbin/models"), objs=[obj], # only this obj select_objs=select_objs, # selected objects - ann_file=osp.join(DATASETS_ROOT, "BOP_DATASETS/ycbv/test_targets_bop19.json"), + ann_file=osp.join(DATASETS_ROOT, "BOP_DATASETS/fruitbin/test_targets_bop19.json"), scale_to_meter=0.001, with_masks=True, # (load masks but may not use it) with_depth=True, # (load depth path here, but may not use it) height=480, width=640, cache_dir=osp.join(PROJ_ROOT, ".cache"), - use_cache=True, + use_cache=False, num_to_load=-1, filter_invalid=False, - ref_key="ycbv", + ref_key="fruitbin", ) @@ -360,25 +360,25 @@ def register_with_name_cfg(name, data_cfg=None): data_cfg can be set in cfg.DATA_CFG.name """ dprint("register dataset: {}".format(name)) - if name in SPLITS_YCBV: - used_cfg = SPLITS_YCBV[name] + if name in SPLITS_FRUITBIN: + used_cfg = SPLITS_FRUITBIN[name] else: assert data_cfg is not None, f"dataset name {name} is not registered" used_cfg = data_cfg - DatasetCatalog.register(name, YCBV_BOP_TEST_Dataset(used_cfg)) + DatasetCatalog.register(name, FRUITBIN_BOP_TEST_Dataset(used_cfg)) # something like eval_types MetadataCatalog.get(name).set( - id="ycbv", # NOTE: for pvnet to determine module + id="fruitbin", # NOTE: for pvnet to determine module ref_key=used_cfg["ref_key"], objs=used_cfg["objs"], eval_error_types=["ad", "rete", "proj"], evaluator_type="bop", - **get_ycbv_metadata(obj_names=used_cfg["objs"], ref_key=used_cfg["ref_key"]), + **get_fruitbin_metadata(obj_names=used_cfg["objs"], ref_key=used_cfg["ref_key"]), ) def get_available_datasets(): - return list(SPLITS_YCBV.keys()) + return list(SPLITS_FRUITBIN.keys()) #### tests ############################################### @@ -440,7 +440,7 @@ if __name__ == "__main__": """Test the dataset loader. Usage: - python -m core.datasets.ycbv_bop_test dataset_name + python -m core.datasets.fruitbin_bop_test dataset_name """ from lib.vis_utils.image import grid_show from lib.utils.setup_logger import setup_my_logger diff --git a/core/gdrn_modeling/datasets/fruitbin_d2.py b/core/gdrn_modeling/datasets/fruitbin_d2.py index 385f463..45d8d35 100755 --- a/core/gdrn_modeling/datasets/fruitbin_d2.py +++ b/core/gdrn_modeling/datasets/fruitbin_d2.py @@ -26,7 +26,7 @@ logger = logging.getLogger(__name__) DATASETS_ROOT = osp.normpath(osp.join(PROJ_ROOT, "datasets")) -class YCBV_Dataset: +class FRUITBIN_Dataset: """use image_sets(scene/image_id) and image root to get data; Here we use bop models, which are center aligned and have some offsets compared to original models.""" @@ -45,9 +45,9 @@ class YCBV_Dataset: self.ann_files = data_cfg["ann_files"] # provide scene/im_id list self.image_prefixes = data_cfg["image_prefixes"] # image root - self.dataset_root = data_cfg["dataset_root"] # BOP_DATASETS/ycbv/ + self.dataset_root = data_cfg["dataset_root"] # BOP_DATASETS/fruitbin/ assert osp.exists(self.dataset_root), self.dataset_root - self.models_root = data_cfg["models_root"] # BOP_DATASETS/ycbv/models + self.models_root = data_cfg["models_root"] # BOP_DATASETS/fruitbin/models self.scale_to_meter = data_cfg["scale_to_meter"] # 0.001 self.with_masks = data_cfg["with_masks"] # True (load masks but may not use it) @@ -77,7 +77,7 @@ class YCBV_Dataset: ################################################## # NOTE: careful! Only the selected objects - self.cat_ids = [cat_id for cat_id, obj_name in ref.ycbv.id2obj.items() if obj_name in self.objs] + self.cat_ids = [cat_id for cat_id, obj_name in ref.fruitbin.id2obj.items() if obj_name in self.objs] # map selected objs to [0, num_objs-1] self.cat2label = {v: i for i, v in enumerate(self.cat_ids)} # id_map self.label2cat = {label: cat for cat, label in self.cat2label.items()} @@ -125,13 +125,13 @@ class YCBV_Dataset: num_instances_without_valid_box = 0 for (scene_id, im_id) in tqdm(scene_im_ids): - rgb_path = osp.join(image_root, f"{scene_id:06d}/rgb/{im_id:06d}.png") + rgb_path = osp.join(image_root, f"{scene_id:06d}/rgb/{im_id:d}.png") assert osp.exists(rgb_path), rgb_path str_im_id = str(im_id) scene_im_id = f"{scene_id}/{im_id}" - # for ycbv/tless, load cam K from image infos + # for fruitbin/tless, load cam K from image infos cam_anno = np.array(scene_cam_dicts[scene_id][str_im_id]["cam_K"], dtype=np.float32).reshape(3, 3) adapth_this_K = False if self.align_K_by_change_pose: @@ -186,7 +186,7 @@ class YCBV_Dataset: quat = mat2quat(pose[:3, :3]) ############# bbox ############################ - bbox = info["bbox_obj"] + bbox = info["bbox_visib"] x1, y1, w, h = bbox x2 = x1 + w y2 = y1 + h @@ -206,7 +206,7 @@ class YCBV_Dataset: if self.with_masks: # either list[list[float]] or dict(RLE) mask_visib_file = osp.join( image_root, - f"{scene_id:06d}/mask_visib/{im_id:06d}_{anno_i:06d}.png", + f"{scene_id:06d}/mask_visib/{im_id:d}.png", ) assert osp.exists(mask_visib_file), mask_visib_file mask = mmcv.imread(mask_visib_file, "unchanged") @@ -216,16 +216,16 @@ class YCBV_Dataset: continue mask_rle = binary_mask_to_rle(mask) - mask_full_file = osp.join( - image_root, - f"{scene_id:06d}/mask/{im_id:06d}_{anno_i:06d}.png", - ) - assert osp.exists(mask_full_file), mask_full_file + # mask_full_file = osp.join( + # image_root, + # f"{scene_id:06d}/mask/{im_id:06d}_{anno_i:06d}.png", + # ) + # assert osp.exists(mask_full_file), mask_full_file # load mask full - mask_full = mmcv.imread(mask_full_file, "unchanged") - mask_full = mask_full.astype("bool") - mask_full_rle = binary_mask_to_rle(mask_full, compressed=True) + # mask_full = mmcv.imread(mask_full_file, "unchanged") + # mask_full = mask_full.astype("bool") + # mask_full_rle = binary_mask_to_rle(mask_full, compressed=True) proj = (self.cam @ trans.T).T # NOTE: use self.cam here proj = proj[:2] / proj[2] @@ -233,14 +233,14 @@ class YCBV_Dataset: inst = { "category_id": cur_label, # 0-based label "bbox": bbox, - "bbox_obj": bbox, + "bbox_visib": bbox, "bbox_mode": BoxMode.XYXY_ABS, "pose": pose, "quat": quat, "trans": trans, "centroid_2d": proj, # absolute (cx, cy) "segmentation": mask_rle, - "mask_full": mask_full_rle, + # "mask_full": mask_full_rle, } if self.with_xyz: @@ -277,7 +277,7 @@ class YCBV_Dataset: ) return dataset_dicts - def __call__(self): # YCBV_Dataset + def __call__(self): # FRUITBIN_Dataset """Load light-weight instance annotations of all images into a list of dicts in Detectron2 format. @@ -346,7 +346,7 @@ class YCBV_Dataset: model = inout.load_ply( osp.join( self.models_root, - f"obj_{ref.ycbv.obj2id[obj_name]:06d}.ply", + f"obj_{ref.fruitbin.obj2id[obj_name]:06d}.ply", ), vertex_scale=self.scale_to_meter, ) @@ -366,7 +366,7 @@ class YCBV_Dataset: ########### register datasets ############################################################ -def get_ycbv_metadata(obj_names, ref_key): +def get_fruitbin_metadata(obj_names, ref_key): """task specific metadata.""" data_ref = ref.__dict__[ref_key] @@ -387,16 +387,16 @@ def get_ycbv_metadata(obj_names, ref_key): return meta -ycbv_model_root = "BOP_DATASETS/ycbv/models/" +fruitbin_model_root = "BOP_DATASETS/fruitbin/models/" ################################################################################ default_cfg = dict( - # name="ycbv_train_real", - dataset_root=osp.join(DATASETS_ROOT, "BOP_DATASETS/ycbv/"), - models_root=osp.join(DATASETS_ROOT, "BOP_DATASETS/ycbv/models"), # models_simple - objs=ref.ycbv.objects, # all objects + # name="fruitbin_train_real", + dataset_root=osp.join(DATASETS_ROOT, "BOP_DATASETS/fruitbin/"), + models_root=osp.join(DATASETS_ROOT, "BOP_DATASETS/fruitbin/models"), # models_simple + objs=ref.fruitbin.objects, # all objects # NOTE: this contains all classes - # ann_files=[osp.join(DATASETS_ROOT, "BOP_DATASETS/ycbv/image_sets/train.txt")], - # image_prefixes=[osp.join(DATASETS_ROOT, "BOP_DATASETS/ycbv/train_real")], + # ann_files=[osp.join(DATASETS_ROOT, "BOP_DATASETS/fruitbin/image_sets/train.txt")], + # image_prefixes=[osp.join(DATASETS_ROOT, "BOP_DATASETS/fruitbin/train_real")], scale_to_meter=0.001, with_masks=True, # (load masks but may not use it) with_depth=True, # (load depth path here, but may not use it) @@ -408,85 +408,85 @@ default_cfg = dict( use_cache=True, num_to_load=-1, filter_invalid=True, - ref_key="ycbv", + ref_key="fruitbin", ) -SPLITS_YCBV = {} +SPLITS_FRUITBIN = {} update_cfgs = { - "ycbv_train_real": { - "ann_files": [osp.join(DATASETS_ROOT, "BOP_DATASETS/ycbv/image_sets/train.txt")], - "image_prefixes": [osp.join(DATASETS_ROOT, "BOP_DATASETS/ycbv/train_real")], + "fruitbin_train_real": { + "ann_files": [osp.join(DATASETS_ROOT, "BOP_DATASETS/fruitbin/image_sets/train.txt")], + "image_prefixes": [osp.join(DATASETS_ROOT, "BOP_DATASETS/fruitbin/train_real")], }, - "ycbv_train_real_aligned_Kuw": { - "ann_files": [osp.join(DATASETS_ROOT, "BOP_DATASETS/ycbv/image_sets/train.txt")], - "image_prefixes": [osp.join(DATASETS_ROOT, "BOP_DATASETS/ycbv/train_real")], + "fruitbin_train_real_aligned_Kuw": { + "ann_files": [osp.join(DATASETS_ROOT, "BOP_DATASETS/fruitbin/image_sets/train.txt")], + "image_prefixes": [osp.join(DATASETS_ROOT, "BOP_DATASETS/fruitbin/train_real")], "align_K_by_change_pose": True, }, - "ycbv_train_real_uw": { - "ann_files": [osp.join(DATASETS_ROOT, "BOP_DATASETS/ycbv/image_sets/train_real_uw.txt")], - "image_prefixes": [osp.join(DATASETS_ROOT, "BOP_DATASETS/ycbv/train_real")], + "fruitbin_train_real_uw": { + "ann_files": [osp.join(DATASETS_ROOT, "BOP_DATASETS/fruitbin/image_sets/train_real_uw.txt")], + "image_prefixes": [osp.join(DATASETS_ROOT, "BOP_DATASETS/fruitbin/train_real")], }, - "ycbv_train_real_uw_every10": { + "fruitbin_train_real_uw_every10": { "ann_files": [ osp.join( DATASETS_ROOT, - "BOP_DATASETS/ycbv/image_sets/train_real_uw_every10.txt", + "BOP_DATASETS/fruitbin/image_sets/train_real_uw_every10.txt", ) ], - "image_prefixes": [osp.join(DATASETS_ROOT, "BOP_DATASETS/ycbv/train_real")], + "image_prefixes": [osp.join(DATASETS_ROOT, "BOP_DATASETS/fruitbin/train_real")], }, - "ycbv_train_real_cmu": { + "fruitbin_train_real_cmu": { "ann_files": [ osp.join( DATASETS_ROOT, - "BOP_DATASETS/ycbv/image_sets/train_real_cmu.txt", + "BOP_DATASETS/fruitbin/image_sets/train_real_cmu.txt", ) ], - "image_prefixes": [osp.join(DATASETS_ROOT, "BOP_DATASETS/ycbv/train_real")], + "image_prefixes": [osp.join(DATASETS_ROOT, "BOP_DATASETS/fruitbin/train_real")], }, - "ycbv_train_real_cmu_aligned_Kuw": { + "fruitbin_train_real_cmu_aligned_Kuw": { "ann_files": [ osp.join( DATASETS_ROOT, - "BOP_DATASETS/ycbv/image_sets/train_real_cmu.txt", + "BOP_DATASETS/fruitbin/image_sets/train_real_cmu.txt", ) ], - "image_prefixes": [osp.join(DATASETS_ROOT, "BOP_DATASETS/ycbv/train_real")], + "image_prefixes": [osp.join(DATASETS_ROOT, "BOP_DATASETS/fruitbin/train_real")], "align_K_by_change_pose": True, }, - "ycbv_train_synt": { - "ann_files": [osp.join(DATASETS_ROOT, "BOP_DATASETS/ycbv/image_sets/train_synt.txt")], - "image_prefixes": [osp.join(DATASETS_ROOT, "BOP_DATASETS/ycbv/train_synt")], + "fruitbin_train_synt": { + "ann_files": [osp.join(DATASETS_ROOT, "BOP_DATASETS/fruitbin/image_sets/train_synt.txt")], + "image_prefixes": [osp.join(DATASETS_ROOT, "BOP_DATASETS/fruitbin/train_synt")], }, - "ycbv_train_synt_50k": { + "fruitbin_train_synt_50k": { "ann_files": [ osp.join( DATASETS_ROOT, - "BOP_DATASETS/ycbv/image_sets/train_synt_50k.txt", + "BOP_DATASETS/fruitbin/image_sets/train_synt_50k.txt", ) ], - "image_prefixes": [osp.join(DATASETS_ROOT, "BOP_DATASETS/ycbv/train_synt")], + "image_prefixes": [osp.join(DATASETS_ROOT, "BOP_DATASETS/fruitbin/train_synt")], }, - "ycbv_train_synt_30k": { + "fruitbin_train_synt_30k": { "ann_files": [ osp.join( DATASETS_ROOT, - "BOP_DATASETS/ycbv/image_sets/train_synt_30k.txt", + "BOP_DATASETS/fruitbin/image_sets/train_synt_30k.txt", ) ], - "image_prefixes": [osp.join(DATASETS_ROOT, "BOP_DATASETS/ycbv/train_synt")], + "image_prefixes": [osp.join(DATASETS_ROOT, "BOP_DATASETS/fruitbin/train_synt")], }, - "ycbv_train_synt_100": { + "fruitbin_train_synt_100": { "ann_files": [ osp.join( DATASETS_ROOT, - "BOP_DATASETS/ycbv/image_sets/train_synt_100.txt", + "BOP_DATASETS/fruitbin/image_sets/train_synt_100.txt", ) ], - "image_prefixes": [osp.join(DATASETS_ROOT, "BOP_DATASETS/ycbv/train_synt")], + "image_prefixes": [osp.join(DATASETS_ROOT, "BOP_DATASETS/fruitbin/train_synt")], }, - "ycbv_test": { - "ann_files": [osp.join(DATASETS_ROOT, "BOP_DATASETS/ycbv/image_sets/keyframe.txt")], - "image_prefixes": [osp.join(DATASETS_ROOT, "BOP_DATASETS/ycbv/test")], + "fruitbin_test": { + "ann_files": [osp.join(DATASETS_ROOT, "BOP_DATASETS/fruitbin/image_sets/keyframe.txt")], + "image_prefixes": [osp.join(DATASETS_ROOT, "BOP_DATASETS/fruitbin/test")], "with_xyz": False, "filter_invalid": False, }, @@ -499,10 +499,10 @@ for name, update_cfg in update_cfgs.items(): if "_100" in name: num_to_load = 100 used_cfg["num_to_load"] = num_to_load - SPLITS_YCBV[name] = used_cfg + SPLITS_FRUITBIN[name] = used_cfg # single object splits ###################################################### -for obj in ref.ycbv.objects: +for obj in ref.fruitbin.objects: for split in [ "train_real", "train_real_aligned_Kuw", @@ -514,7 +514,7 @@ for obj in ref.ycbv.objects: "train_synt_30k", "test", ]: - name = "ycbv_{}_{}".format(obj, split) + name = "fruitbin_{}_{}".format(obj, split) if split in [ "train_real", "train_real_aligned_Kuw", @@ -554,19 +554,19 @@ for obj in ref.ycbv.objects: } root_name, idx_file = split_idx_file_dict[split] - if name not in SPLITS_YCBV: - SPLITS_YCBV[name] = dict( + if name not in SPLITS_FRUITBIN: + SPLITS_FRUITBIN[name] = dict( name=name, - dataset_root=osp.join(DATASETS_ROOT, "BOP_DATASETS/ycbv/"), - models_root=osp.join(DATASETS_ROOT, "BOP_DATASETS/ycbv/models"), + dataset_root=osp.join(DATASETS_ROOT, "BOP_DATASETS/fruitbin/"), + models_root=osp.join(DATASETS_ROOT, "BOP_DATASETS/fruitbin/models"), objs=[obj], ann_files=[ osp.join( DATASETS_ROOT, - "BOP_DATASETS/ycbv/image_sets/{}".format(idx_file), + "BOP_DATASETS/fruitbin/image_sets/{}".format(idx_file), ) ], - image_prefixes=[osp.join(DATASETS_ROOT, "BOP_DATASETS/ycbv/{}".format(root_name))], + image_prefixes=[osp.join(DATASETS_ROOT, "BOP_DATASETS/fruitbin/{}".format(root_name))], scale_to_meter=0.001, with_masks=True, # (load masks but may not use it) with_depth=True, # (load depth path here, but may not use it) @@ -578,7 +578,7 @@ for obj in ref.ycbv.objects: use_cache=True, num_to_load=-1, filter_invalid=filter_invalid, - ref_key="ycbv", + ref_key="fruitbin", ) @@ -592,32 +592,32 @@ def register_with_name_cfg(name, data_cfg=None): data_cfg can be set in cfg.DATA_CFG.name """ dprint("register dataset: {}".format(name)) - if name in SPLITS_YCBV: - used_cfg = SPLITS_YCBV[name] + if name in SPLITS_FRUITBIN: + used_cfg = SPLITS_FRUITBIN[name] else: assert ( data_cfg is not None - ), f"dataset name {name} is not registered. available datasets: {list(SPLITS_YCBV.keys())}" + ), f"dataset name {name} is not registered. available datasets: {list(SPLITS_FRUITBIN.keys())}" used_cfg = data_cfg - DatasetCatalog.register(name, YCBV_Dataset(used_cfg)) + DatasetCatalog.register(name, FRUITBIN_Dataset(used_cfg)) # something like eval_types MetadataCatalog.get(name).set( - id="ycbv", # NOTE: for pvnet to determine module + id="fruitbin", # NOTE: for pvnet to determine module ref_key=used_cfg["ref_key"], objs=used_cfg["objs"], eval_error_types=["ad", "rete", "proj"], evaluator_type="bop", - **get_ycbv_metadata(obj_names=used_cfg["objs"], ref_key=used_cfg["ref_key"]), + **get_fruitbin_metadata(obj_names=used_cfg["objs"], ref_key=used_cfg["ref_key"]), ) def get_available_datasets(): - return list(SPLITS_YCBV.keys()) + return list(SPLITS_FRUITBIN.keys()) #### tests ############################################### def test_vis(): - # python -m core.datasets.ycbv_d2 ycbv_test + # python -m core.datasets.fruitbin_d2 fruitbin_test dataset_name = sys.argv[1] meta = MetadataCatalog.get(dataset_name) t_start = time.perf_counter() @@ -625,7 +625,7 @@ def test_vis(): with_xyz = False if "test" in dataset_name else True logger.info("Done loading {} samples with {:.3f}s.".format(len(dicts), time.perf_counter() - t_start)) - dirname = "output/ycbv_test-data-vis" + dirname = "output/fruitbin_test-data-vis" os.makedirs(dirname, exist_ok=True) objs = meta.objs for d in dicts: @@ -739,3 +739,4 @@ if __name__ == "__main__": register_with_name_cfg(sys.argv[1]) print("dataset catalog: ", DatasetCatalog.list()) test_vis() + diff --git a/core/gdrn_modeling/datasets/fruitbin_pbr.py b/core/gdrn_modeling/datasets/fruitbin_pbr.py index de7ecb2..4cede9e 100644 --- a/core/gdrn_modeling/datasets/fruitbin_pbr.py +++ b/core/gdrn_modeling/datasets/fruitbin_pbr.py @@ -25,7 +25,7 @@ logger = logging.getLogger(__name__) DATASETS_ROOT = osp.normpath(osp.join(PROJ_ROOT, "datasets")) -class YCBV_PBR_Dataset: +class FRUITBIN_PBR_Dataset: def __init__(self, data_cfg): """ Set with_depth and with_masks default to True, @@ -39,7 +39,7 @@ class YCBV_PBR_Dataset: self.dataset_root = data_cfg.get( "dataset_root", - osp.join(DATASETS_ROOT, "BOP_DATASETS/ycbv/train_pbr"), + osp.join(DATASETS_ROOT, "BOP_DATASETS/fruitbin/train_pbr"), ) self.xyz_root = data_cfg.get("xyz_root", osp.join(self.dataset_root, "xyz_crop")) assert osp.exists(self.dataset_root), self.dataset_root @@ -59,14 +59,14 @@ class YCBV_PBR_Dataset: ################################################## # NOTE: careful! Only the selected objects - self.cat_ids = [cat_id for cat_id, obj_name in ref.ycbv.id2obj.items() if obj_name in self.objs] + self.cat_ids = [cat_id for cat_id, obj_name in ref.fruitbin.id2obj.items() if obj_name in self.objs] # map selected objs to [0, num_objs-1] self.cat2label = {v: i for i, v in enumerate(self.cat_ids)} # id_map self.label2cat = {label: cat for cat, label in self.cat2label.items()} self.obj2label = OrderedDict((obj, obj_id) for obj_id, obj in enumerate(self.objs)) ########################################################## - self.scenes = [f"{i:06d}" for i in range(50)] + self.scenes = [f"{i:06d}" for i in range(8)] def __call__(self): """Load light-weight instance annotations of all images into a list of @@ -111,13 +111,12 @@ class YCBV_PBR_Dataset: gt_dict = mmcv.load(osp.join(scene_root, "scene_gt.json")) gt_info_dict = mmcv.load(osp.join(scene_root, "scene_gt_info.json")) cam_dict = mmcv.load(osp.join(scene_root, "scene_camera.json")) - for str_im_id in tqdm(gt_dict, postfix=f"{scene_id}"): int_im_id = int(str_im_id) - rgb_path = osp.join(scene_root, "rgb/{:06d}.jpg").format(int_im_id) + rgb_path = osp.join(scene_root, "rgb/{:d}.png").format(int_im_id) assert osp.exists(rgb_path), rgb_path - depth_path = osp.join(scene_root, "depth/{:06d}.png".format(int_im_id)) + depth_path = osp.join(scene_root, "depth/{:d}.png".format(int_im_id)) scene_im_id = f"{scene_id}/{int_im_id}" @@ -151,7 +150,7 @@ class YCBV_PBR_Dataset: proj = proj[:2] / proj[2] bbox_visib = gt_info_dict[str_im_id][anno_i]["bbox_visib"] - bbox_obj = gt_info_dict[str_im_id][anno_i]["bbox_obj"] + # bbox_obj = gt_info_dict[str_im_id][anno_i]["bbox_obj"] x1, y1, w, h = bbox_visib if self.filter_invalid: if h <= 1 or w <= 1: @@ -164,9 +163,9 @@ class YCBV_PBR_Dataset: ) mask_visib_file = osp.join( scene_root, - "mask_visib/{:06d}_{:06d}.png".format(int_im_id, anno_i), + "mask_visib/{:d}.png".format(int_im_id), ) - assert osp.exists(mask_file), mask_file + # assert osp.exists(mask_file), mask_file assert osp.exists(mask_visib_file), mask_visib_file # load mask visib TODO: load both mask_visib and mask_full mask_single = mmcv.imread(mask_visib_file, "unchanged") @@ -177,9 +176,9 @@ class YCBV_PBR_Dataset: mask_rle = binary_mask_to_rle(mask_single, compressed=True) # load mask full - mask_full = mmcv.imread(mask_file, "unchanged") - mask_full = mask_full.astype("bool") - mask_full_rle = binary_mask_to_rle(mask_full, compressed=True) + # mask_full = mmcv.imread(mask_file, "unchanged") + # mask_full = mask_full.astype("bool") + # mask_full_rle = binary_mask_to_rle(mask_full, compressed=True) visib_fract = gt_info_dict[str_im_id][anno_i].get("visib_fract", 1.0) @@ -188,14 +187,14 @@ class YCBV_PBR_Dataset: inst = { "category_id": cur_label, # 0-based label "bbox": bbox_visib, # TODO: load both bbox_obj and bbox_visib - "bbox_obj": bbox_obj, + # "bbox_obj": bbox_obj, "bbox_mode": BoxMode.XYWH_ABS, "pose": pose, "quat": quat, "trans": t, "centroid_2d": proj, # absolute (cx, cy) "segmentation": mask_rle, - "mask_full": mask_full_rle, # TODO: load as mask_full, rle + # "mask_full": mask_full_rle, # TODO: load as mask_full, rle "visib_fract": visib_fract, "xyz_path": xyz_path, } @@ -254,7 +253,7 @@ class YCBV_PBR_Dataset: model = inout.load_ply( osp.join( self.models_root, - f"obj_{ref.ycbv.obj2id[obj_name]:06d}.ply", + f"obj_{ref.fruitbin.obj2id[obj_name]:06d}.ply", ), vertex_scale=self.scale_to_meter, ) @@ -274,7 +273,7 @@ class YCBV_PBR_Dataset: ########### register datasets ############################################################ -def get_ycbv_metadata(obj_names, ref_key): +def get_fruitbin_metadata(obj_names, ref_key): """task specific metadata.""" data_ref = ref.__dict__[ref_key] @@ -295,55 +294,55 @@ def get_ycbv_metadata(obj_names, ref_key): return meta -ycbv_model_root = "BOP_DATASETS/ycbv/models/" +fruitbin_model_root = "BOP_DATASETS/fruitbin/models/" ################################################################################ -SPLITS_YCBV_PBR = dict( - ycbv_train_pbr=dict( - name="ycbv_train_pbr", - objs=ref.ycbv.objects, # selected objects - dataset_root=osp.join(DATASETS_ROOT, "BOP_DATASETS/ycbv/train_pbr"), - models_root=osp.join(DATASETS_ROOT, "BOP_DATASETS/ycbv/models"), - xyz_root=osp.join(DATASETS_ROOT, "BOP_DATASETS/ycbv/train_pbr/xyz_crop"), +SPLITS_FRUITBIN_PBR = dict( + fruitbin_train_pbr=dict( + name="fruitbin_train_pbr", + objs=ref.fruitbin.objects, # selected objects + dataset_root=osp.join(DATASETS_ROOT, "BOP_DATASETS/fruitbin/train_pbr"), + models_root=osp.join(DATASETS_ROOT, "BOP_DATASETS/fruitbin/models"), + xyz_root=osp.join(DATASETS_ROOT, "BOP_DATASETS/fruitbin/train_pbr/xyz_crop"), scale_to_meter=0.001, with_masks=True, # (load masks but may not use it) with_depth=True, # (load depth path here, but may not use it) height=480, width=640, - use_cache=True, + use_cache=False, num_to_load=-1, filter_invalid=True, - ref_key="ycbv", + ref_key="fruitbin", ) ) # single obj splits -for obj in ref.ycbv.objects: +for obj in ref.fruitbin.objects: for split in ["train_pbr"]: - name = "ycbv_{}_{}".format(obj, split) + name = "fruitbin_{}_{}".format(obj, split) if split in ["train_pbr"]: filter_invalid = True elif split in ["test"]: filter_invalid = False else: raise ValueError("{}".format(split)) - if name not in SPLITS_YCBV_PBR: - SPLITS_YCBV_PBR[name] = dict( + if name not in SPLITS_FRUITBIN_PBR: + SPLITS_FRUITBIN_PBR[name] = dict( name=name, objs=[obj], # only this obj - dataset_root=osp.join(DATASETS_ROOT, "BOP_DATASETS/ycbv/train_pbr"), - models_root=osp.join(DATASETS_ROOT, "BOP_DATASETS/ycbv/models"), - xyz_root=osp.join(DATASETS_ROOT, "BOP_DATASETS/ycbv/train_pbr/xyz_crop"), + dataset_root=osp.join(DATASETS_ROOT, "BOP_DATASETS/fruitbin/train_pbr"), + models_root=osp.join(DATASETS_ROOT, "BOP_DATASETS/fruitbin/models"), + xyz_root=osp.join(DATASETS_ROOT, "BOP_DATASETS/fruitbin/train_pbr/xyz_crop"), scale_to_meter=0.001, with_masks=True, # (load masks but may not use it) with_depth=True, # (load depth path here, but may not use it) height=480, width=640, - use_cache=True, + use_cache=False, num_to_load=-1, filter_invalid=filter_invalid, - ref_key="ycbv", + ref_key="fruitbin", ) @@ -357,25 +356,25 @@ def register_with_name_cfg(name, data_cfg=None): data_cfg can be set in cfg.DATA_CFG.name """ dprint("register dataset: {}".format(name)) - if name in SPLITS_YCBV_PBR: - used_cfg = SPLITS_YCBV_PBR[name] + if name in SPLITS_FRUITBIN_PBR: + used_cfg = SPLITS_FRUITBIN_PBR[name] else: assert data_cfg is not None, f"dataset name {name} is not registered" used_cfg = data_cfg - DatasetCatalog.register(name, YCBV_PBR_Dataset(used_cfg)) + DatasetCatalog.register(name, FRUITBIN_PBR_Dataset(used_cfg)) # something like eval_types MetadataCatalog.get(name).set( - id="ycbv", # NOTE: for pvnet to determine module + id="fruitbin", # NOTE: for pvnet to determine module ref_key=used_cfg["ref_key"], objs=used_cfg["objs"], eval_error_types=["ad", "rete", "proj"], evaluator_type="bop", - **get_ycbv_metadata(obj_names=used_cfg["objs"], ref_key=used_cfg["ref_key"]), + **get_fruitbin_metadata(obj_names=used_cfg["objs"], ref_key=used_cfg["ref_key"]), ) def get_available_datasets(): - return list(SPLITS_YCBV_PBR.keys()) + return list(SPLITS_FRUITBIN_PBR.keys()) #### tests ############################################### @@ -473,7 +472,7 @@ if __name__ == "__main__": """Test the dataset loader. Usage: - python -m this_module ycbv_pbr_train + python -m this_module fruitbin_pbr_train """ from lib.vis_utils.image import grid_show from lib.utils.setup_logger import setup_my_logger diff --git a/core/gdrn_modeling/main_gdrn.py b/core/gdrn_modeling/main_gdrn.py index 12a5339..4438c12 100644 --- a/core/gdrn_modeling/main_gdrn.py +++ b/core/gdrn_modeling/main_gdrn.py @@ -204,7 +204,6 @@ def main(args): precision=16 if cfg.SOLVER.AMP.ENABLED else 32, ).run(args, cfg) - if __name__ == "__main__": import resource diff --git a/core/gdrn_modeling/tools/fruitbin/convert_det_to_our_format.py b/core/gdrn_modeling/tools/fruitbin/convert_det_to_our_format.py deleted file mode 100644 index e5de9ea..0000000 --- a/core/gdrn_modeling/tools/fruitbin/convert_det_to_our_format.py +++ /dev/null @@ -1,69 +0,0 @@ -import mmcv -import sys -from tqdm import tqdm - -import json - -path = "/data2/lxy/Storage/bop22_results/yolovx_amodal/ycbv/yolox_x_640_ycbv_real_pbr_ycbv_bop_test.json" -ds = mmcv.load(path) - -outs = {} -for d in tqdm(ds): - scene_id = d["scene_id"] - image_id = d["image_id"] - scene_im_id = f"{scene_id}/{image_id}" - - obj_id = d["category_id"] - score = d["score"] - - bbox = d["bbox"] - time = d["time"] - - cur_dict = { - "bbox_est": bbox, - "obj_id": obj_id, - "score": score, - "time": time, - } - - if scene_im_id in outs.keys(): - outs[scene_im_id].append(cur_dict) - else: - outs[scene_im_id] = [cur_dict] - - -def save_json(path, content, sort=False): - """Saves the provided content to a JSON file. - - :param path: Path to the output JSON file. - :param content: Dictionary/list to save. - """ - with open(path, "w") as f: - - if isinstance(content, dict): - f.write("{\n") - if sort: - content_sorted = sorted(content.items(), key=lambda x: x[0]) - else: - content_sorted = content.items() - for elem_id, (k, v) in enumerate(content_sorted): - f.write(' "{}": {}'.format(k, json.dumps(v, sort_keys=True))) - if elem_id != len(content) - 1: - f.write(",") - f.write("\n") - f.write("}") - - elif isinstance(content, list): - f.write("[\n") - for elem_id, elem in enumerate(content): - f.write(" {}".format(json.dumps(elem, sort_keys=True))) - if elem_id != len(content) - 1: - f.write(",") - f.write("\n") - f.write("]") - - else: - json.dump(content, f, sort_keys=True) - - -save_json("datasets/BOP_DATASETS/ycbv/test/test_bboxes/yolox_x_640_ycbv_real_pbr_ycbv_bop_test.json", outs) diff --git a/core/gdrn_modeling/tools/fruitbin/fruitbin_1_compute_fps.py b/core/gdrn_modeling/tools/fruitbin/fruitbin_1_compute_fps.py new file mode 100644 index 0000000..0686824 --- /dev/null +++ b/core/gdrn_modeling/tools/fruitbin/fruitbin_1_compute_fps.py @@ -0,0 +1,45 @@ +# compute fps (farthest point sampling) for models +import os.path as osp +import sys +from tqdm import tqdm + +cur_dir = osp.dirname(osp.abspath(__file__)) +sys.path.insert(0, osp.join(cur_dir, "../../../../")) +import mmcv +from lib.pysixd import inout, misc +import ref +from core.utils.data_utils import get_fps_and_center + + +ref_key = "fruitbin" +data_ref = ref.__dict__[ref_key] + +model_dir = data_ref.model_dir +id2obj = data_ref.id2obj + + +def main(): + vertex_scale = 0.001 + fps_dict = {} + for obj_id in tqdm(id2obj): + print(obj_id) + model_path = osp.join(model_dir, f"obj_{obj_id:06d}.ply") + model = inout.load_ply(model_path, vertex_scale=vertex_scale) + fps_dict[str(obj_id)] = {} + fps_dict[str(obj_id)]["fps4_and_center"] = get_fps_and_center(model["pts"], num_fps=4, init_center=True) + fps_dict[str(obj_id)]["fps8_and_center"] = get_fps_and_center(model["pts"], num_fps=8, init_center=True) + fps_dict[str(obj_id)]["fps12_and_center"] = get_fps_and_center(model["pts"], num_fps=12, init_center=True) + fps_dict[str(obj_id)]["fps16_and_center"] = get_fps_and_center(model["pts"], num_fps=16, init_center=True) + fps_dict[str(obj_id)]["fps20_and_center"] = get_fps_and_center(model["pts"], num_fps=20, init_center=True) + fps_dict[str(obj_id)]["fps32_and_center"] = get_fps_and_center(model["pts"], num_fps=32, init_center=True) + fps_dict[str(obj_id)]["fps64_and_center"] = get_fps_and_center(model["pts"], num_fps=64, init_center=True) + fps_dict[str(obj_id)]["fps128_and_center"] = get_fps_and_center(model["pts"], num_fps=128, init_center=True) + fps_dict[str(obj_id)]["fps256_and_center"] = get_fps_and_center(model["pts"], num_fps=256, init_center=True) + + save_path = osp.join(model_dir, "fps_points.pkl") + mmcv.dump(fps_dict, save_path) + print(f"saved to {save_path}") + + +if __name__ == "__main__": + main() diff --git a/det/yolox/data/datasets/dataset_factory.py b/det/yolox/data/datasets/dataset_factory.py index eabeef9..c83393b 100644 --- a/det/yolox/data/datasets/dataset_factory.py +++ b/det/yolox/data/datasets/dataset_factory.py @@ -22,6 +22,9 @@ from . import ( itodd_pbr, itodd_d2, itodd_bop_test, + fruitbin_pbr, + fruitbin_d2, + fruitbin_bop_test, ) # noqa cur_dir = osp.dirname(osp.abspath(__file__)) @@ -53,6 +56,9 @@ _DSET_MOD_NAMES = [ "itodd_pbr", "itodd_d2", "itodd_bop_test", + "fruitbin_pbr", + "fruitbin_d2", + "fruitbin_bop_test", ] logger = logging.getLogger(__name__) diff --git a/det/yolox/data/datasets/fruitbin_bop_test.py b/det/yolox/data/datasets/fruitbin_bop_test.py index c9a9324..fd1abeb 100644 --- a/det/yolox/data/datasets/fruitbin_bop_test.py +++ b/det/yolox/data/datasets/fruitbin_bop_test.py @@ -25,8 +25,8 @@ logger = logging.getLogger(__name__) DATASETS_ROOT = osp.normpath(osp.join(PROJ_ROOT, "datasets")) -class YCBV_BOP_TEST_Dataset: - """ycbv bop test.""" +class FRUITBIN_BOP_TEST_Dataset: + """fruitbin bop test.""" def __init__(self, data_cfg): """ @@ -43,8 +43,8 @@ class YCBV_BOP_TEST_Dataset: self.ann_file = data_cfg["ann_file"] # json file with scene_id and im_id items - self.dataset_root = data_cfg["dataset_root"] # BOP_DATASETS/ycbv/test - self.models_root = data_cfg["models_root"] # BOP_DATASETS/ycbv/models + self.dataset_root = data_cfg["dataset_root"] # BOP_DATASETS/fruitbin/test + self.models_root = data_cfg["models_root"] # BOP_DATASETS/fruitbin/models self.scale_to_meter = data_cfg["scale_to_meter"] # 0.001 self.with_masks = data_cfg["with_masks"] # True (load masks but may not use it) @@ -60,7 +60,7 @@ class YCBV_BOP_TEST_Dataset: ################################################## # NOTE: careful! Only the selected objects - self.cat_ids = [cat_id for cat_id, obj_name in ref.ycbv.id2obj.items() if obj_name in self.objs] + self.cat_ids = [cat_id for cat_id, obj_name in ref.fruitbin.id2obj.items() if obj_name in self.objs] # map selected objs to [0, num_objs-1] self.cat2label = {v: i for i, v in enumerate(self.cat_ids)} # id_map self.label2cat = {label: cat for cat, label in self.cat2label.items()} @@ -127,10 +127,10 @@ class YCBV_BOP_TEST_Dataset: for scene_id, im_id in tqdm(scene_im_ids): str_im_id = str(im_id) scene_root = osp.join(self.dataset_root, f"{scene_id:06d}") - rgb_path = osp.join(scene_root, "rgb/{:06d}.png").format(im_id) + rgb_path = osp.join(scene_root, "rgb/{:d}.png").format(im_id) assert osp.exists(rgb_path), rgb_path - depth_path = osp.join(scene_root, "depth/{:06d}.png".format(im_id)) + depth_path = osp.join(scene_root, "depth/{:d}.png".format(im_id)) scene_id = int(rgb_path.split("/")[-3]) @@ -152,7 +152,7 @@ class YCBV_BOP_TEST_Dataset: insts = [] for anno_i, anno in enumerate(gt_dicts[scene_id][str_im_id]): obj_id = anno["obj_id"] - if ref.ycbv.id2obj[obj_id] not in self.select_objs: + if ref.fruitbin.id2obj[obj_id] not in self.select_objs: continue cur_label = self.cat2label[obj_id] # 0-based label R = np.array(anno["cam_R_m2c"], dtype="float32").reshape(3, 3) @@ -164,22 +164,22 @@ class YCBV_BOP_TEST_Dataset: proj = proj[:2] / proj[2] bbox_visib = gt_info_dicts[scene_id][str_im_id][anno_i]["bbox_visib"] - bbox_obj = gt_info_dicts[scene_id][str_im_id][anno_i]["bbox_obj"] + # bbox_obj = gt_info_dicts[scene_id][str_im_id][anno_i]["bbox_obj"] x1, y1, w, h = bbox_visib if self.filter_invalid: if h <= 1 or w <= 1: self.num_instances_without_valid_box += 1 continue - mask_file = osp.join( - scene_root, - "mask/{:06d}_{:06d}.png".format(im_id, anno_i), - ) + # mask_file = osp.join( + # scene_root, + # "mask/{:06d}_{:06d}.png".format(im_id, anno_i), + # ) mask_visib_file = osp.join( scene_root, - "mask_visib/{:06d}_{:06d}.png".format(im_id, anno_i), + "mask_visib/{:d}.png".format(im_id, anno_i), ) - assert osp.exists(mask_file), mask_file + # assert osp.exists(mask_file), mask_file assert osp.exists(mask_visib_file), mask_visib_file # load mask visib mask_single = mmcv.imread(mask_visib_file, "unchanged") @@ -190,9 +190,9 @@ class YCBV_BOP_TEST_Dataset: mask_rle = binary_mask_to_rle(mask_single, compressed=True) # load mask full - mask_full = mmcv.imread(mask_file, "unchanged") - mask_full = mask_full.astype("bool") - mask_full_rle = binary_mask_to_rle(mask_full, compressed=True) + # mask_full = mmcv.imread(mask_file, "unchanged") + # mask_full = mask_full.astype("bool") + # mask_full_rle = binary_mask_to_rle(mask_full, compressed=True) inst = { "category_id": cur_label, # 0-based label @@ -203,7 +203,7 @@ class YCBV_BOP_TEST_Dataset: "trans": t, "centroid_2d": proj, # absolute (cx, cy) "segmentation": mask_rle, - "mask_full": mask_full_rle, # TODO: load as mask_full, rle + # "mask_full": mask_full_rle, # TODO: load as mask_full, rle } model_info = self.models_info[str(obj_id)] @@ -260,7 +260,7 @@ class YCBV_BOP_TEST_Dataset: model = inout.load_ply( osp.join( self.models_root, - f"obj_{ref.ycbv.obj2id[obj_name]:06d}.ply", + f"obj_{ref.fruitbin.obj2id[obj_name]:06d}.ply", ), vertex_scale=self.scale_to_meter, ) @@ -280,7 +280,7 @@ class YCBV_BOP_TEST_Dataset: ########### register datasets ############################################################ -def get_ycbv_metadata(obj_names, ref_key): +def get_fruitbin_metadata(obj_names, ref_key): """task specific metadata.""" data_ref = ref.__dict__[ref_key] @@ -303,49 +303,49 @@ def get_ycbv_metadata(obj_names, ref_key): ################################################################################ -SPLITS_YCBV = dict( - ycbv_bop_test=dict( - name="ycbv_bop_test", - dataset_root=osp.join(DATASETS_ROOT, "BOP_DATASETS/ycbv/test"), - models_root=osp.join(DATASETS_ROOT, "BOP_DATASETS/ycbv/models"), - objs=ref.ycbv.objects, # selected objects - ann_file=osp.join(DATASETS_ROOT, "BOP_DATASETS/ycbv/test_targets_bop19.json"), +SPLITS_FRUITBIN = dict( + fruitbin_bop_test=dict( + name="fruitbin_bop_test", + dataset_root=osp.join(DATASETS_ROOT, "BOP_DATASETS/fruitbin/test"), + models_root=osp.join(DATASETS_ROOT, "BOP_DATASETS/fruitbin/models"), + objs=ref.fruitbin.objects, # selected objects + ann_file=osp.join(DATASETS_ROOT, "BOP_DATASETS/fruitbin/test_targets_bop19.json"), scale_to_meter=0.001, with_masks=True, # (load masks but may not use it) with_depth=True, # (load depth path here, but may not use it) height=480, width=640, cache_dir=osp.join(PROJ_ROOT, ".cache"), - use_cache=True, + use_cache=False, num_to_load=-1, filter_invalid=False, - ref_key="ycbv", + ref_key="fruitbin", ) ) # single objs (num_class is from all objs) -for obj in ref.ycbv.objects: - name = "ycbv_bop_{}_test".format(obj) +for obj in ref.fruitbin.objects: + name = "fruitbin_bop_{}_test".format(obj) select_objs = [obj] - if name not in SPLITS_YCBV: - SPLITS_YCBV[name] = dict( + if name not in SPLITS_FRUITBIN: + SPLITS_FRUITBIN[name] = dict( name=name, - dataset_root=osp.join(DATASETS_ROOT, "BOP_DATASETS/ycbv/test"), - models_root=osp.join(DATASETS_ROOT, "BOP_DATASETS/ycbv/models"), + dataset_root=osp.join(DATASETS_ROOT, "BOP_DATASETS/fruitbin/test"), + models_root=osp.join(DATASETS_ROOT, "BOP_DATASETS/fruitbin/models"), objs=[obj], # only this obj select_objs=select_objs, # selected objects - ann_file=osp.join(DATASETS_ROOT, "BOP_DATASETS/ycbv/test_targets_bop19.json"), + ann_file=osp.join(DATASETS_ROOT, "BOP_DATASETS/fruitbin/test_targets_bop19.json"), scale_to_meter=0.001, with_masks=True, # (load masks but may not use it) with_depth=True, # (load depth path here, but may not use it) height=480, width=640, cache_dir=osp.join(PROJ_ROOT, ".cache"), - use_cache=True, + use_cache=False, num_to_load=-1, filter_invalid=False, - ref_key="ycbv", + ref_key="fruitbin", ) @@ -359,25 +359,25 @@ def register_with_name_cfg(name, data_cfg=None): data_cfg can be set in cfg.DATA_CFG.name """ dprint("register dataset: {}".format(name)) - if name in SPLITS_YCBV: - used_cfg = SPLITS_YCBV[name] + if name in SPLITS_FRUITBIN: + used_cfg = SPLITS_FRUITBIN[name] else: assert data_cfg is not None, f"dataset name {name} is not registered" used_cfg = data_cfg - DatasetCatalog.register(name, YCBV_BOP_TEST_Dataset(used_cfg)) + DatasetCatalog.register(name, FRUITBIN_BOP_TEST_Dataset(used_cfg)) # something like eval_types MetadataCatalog.get(name).set( - id="ycbv", # NOTE: for pvnet to determine module + id="fruitbin", # NOTE: for pvnet to determine module ref_key=used_cfg["ref_key"], objs=used_cfg["objs"], eval_error_types=["ad", "rete", "proj"], evaluator_type="bop", - **get_ycbv_metadata(obj_names=used_cfg["objs"], ref_key=used_cfg["ref_key"]), + **get_fruitbin_metadata(obj_names=used_cfg["objs"], ref_key=used_cfg["ref_key"]), ) def get_available_datasets(): - return list(SPLITS_YCBV.keys()) + return list(SPLITS_FRUITBIN.keys()) #### tests ############################################### @@ -439,7 +439,7 @@ if __name__ == "__main__": """Test the dataset loader. Usage: - python -m core.datasets.ycbv_bop_test dataset_name + python -m core.datasets.fruitbin_bop_test dataset_name """ from lib.vis_utils.image import grid_show from lib.utils.setup_logger import setup_my_logger @@ -451,4 +451,4 @@ if __name__ == "__main__": logger = setup_my_logger(name="core") register_with_name_cfg(sys.argv[1]) print("dataset catalog: ", DatasetCatalog.list()) - test_vis() + test_vis() \ No newline at end of file diff --git a/det/yolox/data/datasets/fruitbin_d2.py b/det/yolox/data/datasets/fruitbin_d2.py index 9534181..d121121 100755 --- a/det/yolox/data/datasets/fruitbin_d2.py +++ b/det/yolox/data/datasets/fruitbin_d2.py @@ -26,7 +26,7 @@ logger = logging.getLogger(__name__) DATASETS_ROOT = osp.normpath(osp.join(PROJ_ROOT, "datasets")) -class YCBV_Dataset: +class FRUITBIN_Dataset: """use image_sets(scene/image_id) and image root to get data; Here we use bop models, which are center aligned and have some offsets compared to original models.""" @@ -45,9 +45,9 @@ class YCBV_Dataset: self.ann_files = data_cfg["ann_files"] # provide scene/im_id list self.image_prefixes = data_cfg["image_prefixes"] # image root - self.dataset_root = data_cfg["dataset_root"] # BOP_DATASETS/ycbv/ + self.dataset_root = data_cfg["dataset_root"] # BOP_DATASETS/fruitbin/ assert osp.exists(self.dataset_root), self.dataset_root - self.models_root = data_cfg["models_root"] # BOP_DATASETS/ycbv/models + self.models_root = data_cfg["models_root"] # BOP_DATASETS/fruitbin/models self.scale_to_meter = data_cfg["scale_to_meter"] # 0.001 self.with_masks = data_cfg["with_masks"] # True (load masks but may not use it) @@ -77,7 +77,7 @@ class YCBV_Dataset: ################################################## # NOTE: careful! Only the selected objects - self.cat_ids = [cat_id for cat_id, obj_name in ref.ycbv.id2obj.items() if obj_name in self.objs] + self.cat_ids = [cat_id for cat_id, obj_name in ref.fruitbin.id2obj.items() if obj_name in self.objs] # map selected objs to [0, num_objs-1] self.cat2label = {v: i for i, v in enumerate(self.cat_ids)} # id_map self.label2cat = {label: cat for cat, label in self.cat2label.items()} @@ -125,13 +125,13 @@ class YCBV_Dataset: num_instances_without_valid_box = 0 for (scene_id, im_id) in tqdm(scene_im_ids): - rgb_path = osp.join(image_root, f"{scene_id:06d}/rgb/{im_id:06d}.png") + rgb_path = osp.join(image_root, f"{scene_id:06d}/rgb/{im_id:d}.png") assert osp.exists(rgb_path), rgb_path str_im_id = str(im_id) scene_im_id = f"{scene_id}/{im_id}" - # for ycbv/tless, load cam K from image infos + # for fruitbin/tless, load cam K from image infos cam_anno = np.array(scene_cam_dicts[scene_id][str_im_id]["cam_K"], dtype=np.float32).reshape(3, 3) adapth_this_K = False if self.align_K_by_change_pose: @@ -186,7 +186,7 @@ class YCBV_Dataset: quat = mat2quat(pose[:3, :3]) ############# bbox ############################ - bbox = info["bbox_obj"] + bbox = info["bbox_visib"] x1, y1, w, h = bbox x2 = x1 + w y2 = y1 + h @@ -206,7 +206,7 @@ class YCBV_Dataset: if self.with_masks: # either list[list[float]] or dict(RLE) mask_visib_file = osp.join( image_root, - f"{scene_id:06d}/mask_visib/{im_id:06d}_{anno_i:06d}.png", + f"{scene_id:06d}/mask_visib/{im_id:d}.png", ) assert osp.exists(mask_visib_file), mask_visib_file mask = mmcv.imread(mask_visib_file, "unchanged") @@ -216,16 +216,16 @@ class YCBV_Dataset: continue mask_rle = binary_mask_to_rle(mask) - mask_full_file = osp.join( - image_root, - f"{scene_id:06d}/mask/{im_id:06d}_{anno_i:06d}.png", - ) - assert osp.exists(mask_full_file), mask_full_file + # mask_full_file = osp.join( + # image_root, + # f"{scene_id:06d}/mask/{im_id:06d}_{anno_i:06d}.png", + # ) + # assert osp.exists(mask_full_file), mask_full_file # load mask full - mask_full = mmcv.imread(mask_full_file, "unchanged") - mask_full = mask_full.astype("bool") - mask_full_rle = binary_mask_to_rle(mask_full, compressed=True) + # mask_full = mmcv.imread(mask_full_file, "unchanged") + # mask_full = mask_full.astype("bool") + # mask_full_rle = binary_mask_to_rle(mask_full, compressed=True) proj = (self.cam @ trans.T).T # NOTE: use self.cam here proj = proj[:2] / proj[2] @@ -239,7 +239,7 @@ class YCBV_Dataset: "trans": trans, "centroid_2d": proj, # absolute (cx, cy) "segmentation": mask_rle, - "mask_full": mask_full_rle, + # "mask_full": mask_full_rle, } if self.with_xyz: @@ -275,7 +275,7 @@ class YCBV_Dataset: ) return dataset_dicts - def __call__(self): # YCBV_Dataset + def __call__(self): # FRUITBIN_Dataset """Load light-weight instance annotations of all images into a list of dicts in Detectron2 format. @@ -344,7 +344,7 @@ class YCBV_Dataset: model = inout.load_ply( osp.join( self.models_root, - f"obj_{ref.ycbv.obj2id[obj_name]:06d}.ply", + f"obj_{ref.fruitbin.obj2id[obj_name]:06d}.ply", ), vertex_scale=self.scale_to_meter, ) @@ -364,7 +364,7 @@ class YCBV_Dataset: ########### register datasets ############################################################ -def get_ycbv_metadata(obj_names, ref_key): +def get_fruitbin_metadata(obj_names, ref_key): """task specific metadata.""" data_ref = ref.__dict__[ref_key] @@ -385,16 +385,16 @@ def get_ycbv_metadata(obj_names, ref_key): return meta -ycbv_model_root = "BOP_DATASETS/ycbv/models/" +fruitbin_model_root = "BOP_DATASETS/fruitbin/models/" ################################################################################ default_cfg = dict( - # name="ycbv_train_real", - dataset_root=osp.join(DATASETS_ROOT, "BOP_DATASETS/ycbv/"), - models_root=osp.join(DATASETS_ROOT, "BOP_DATASETS/ycbv/models"), # models_simple - objs=ref.ycbv.objects, # all objects + # name="fruitbin_train_real", + dataset_root=osp.join(DATASETS_ROOT, "BOP_DATASETS/fruitbin/"), + models_root=osp.join(DATASETS_ROOT, "BOP_DATASETS/fruitbin/models"), # models_simple + objs=ref.fruitbin.objects, # all objects # NOTE: this contains all classes - # ann_files=[osp.join(DATASETS_ROOT, "BOP_DATASETS/ycbv/image_sets/train.txt")], - # image_prefixes=[osp.join(DATASETS_ROOT, "BOP_DATASETS/ycbv/train_real")], + # ann_files=[osp.join(DATASETS_ROOT, "BOP_DATASETS/fruitbin/image_sets/train.txt")], + # image_prefixes=[osp.join(DATASETS_ROOT, "BOP_DATASETS/fruitbin/train_real")], scale_to_meter=0.001, with_masks=True, # (load masks but may not use it) with_depth=True, # (load depth path here, but may not use it) @@ -406,85 +406,85 @@ default_cfg = dict( use_cache=True, num_to_load=-1, filter_invalid=True, - ref_key="ycbv", + ref_key="fruitbin", ) -SPLITS_YCBV = {} +SPLITS_FRUITBIN = {} update_cfgs = { - "ycbv_train_real": { - "ann_files": [osp.join(DATASETS_ROOT, "BOP_DATASETS/ycbv/image_sets/train.txt")], - "image_prefixes": [osp.join(DATASETS_ROOT, "BOP_DATASETS/ycbv/train_real")], + "fruitbin_train_real": { + "ann_files": [osp.join(DATASETS_ROOT, "BOP_DATASETS/fruitbin/image_sets/train.txt")], + "image_prefixes": [osp.join(DATASETS_ROOT, "BOP_DATASETS/fruitbin/train_real")], }, - "ycbv_train_real_aligned_Kuw": { - "ann_files": [osp.join(DATASETS_ROOT, "BOP_DATASETS/ycbv/image_sets/train.txt")], - "image_prefixes": [osp.join(DATASETS_ROOT, "BOP_DATASETS/ycbv/train_real")], + "fruitbin_train_real_aligned_Kuw": { + "ann_files": [osp.join(DATASETS_ROOT, "BOP_DATASETS/fruitbin/image_sets/train.txt")], + "image_prefixes": [osp.join(DATASETS_ROOT, "BOP_DATASETS/fruitbin/train_real")], "align_K_by_change_pose": True, }, - "ycbv_train_real_uw": { - "ann_files": [osp.join(DATASETS_ROOT, "BOP_DATASETS/ycbv/image_sets/train_real_uw.txt")], - "image_prefixes": [osp.join(DATASETS_ROOT, "BOP_DATASETS/ycbv/train_real")], + "fruitbin_train_real_uw": { + "ann_files": [osp.join(DATASETS_ROOT, "BOP_DATASETS/fruitbin/image_sets/train_real_uw.txt")], + "image_prefixes": [osp.join(DATASETS_ROOT, "BOP_DATASETS/fruitbin/train_real")], }, - "ycbv_train_real_uw_every10": { + "fruitbin_train_real_uw_every10": { "ann_files": [ osp.join( DATASETS_ROOT, - "BOP_DATASETS/ycbv/image_sets/train_real_uw_every10.txt", + "BOP_DATASETS/fruitbin/image_sets/train_real_uw_every10.txt", ) ], - "image_prefixes": [osp.join(DATASETS_ROOT, "BOP_DATASETS/ycbv/train_real")], + "image_prefixes": [osp.join(DATASETS_ROOT, "BOP_DATASETS/fruitbin/train_real")], }, - "ycbv_train_real_cmu": { + "fruitbin_train_real_cmu": { "ann_files": [ osp.join( DATASETS_ROOT, - "BOP_DATASETS/ycbv/image_sets/train_real_cmu.txt", + "BOP_DATASETS/fruitbin/image_sets/train_real_cmu.txt", ) ], - "image_prefixes": [osp.join(DATASETS_ROOT, "BOP_DATASETS/ycbv/train_real")], + "image_prefixes": [osp.join(DATASETS_ROOT, "BOP_DATASETS/fruitbin/train_real")], }, - "ycbv_train_real_cmu_aligned_Kuw": { + "fruitbin_train_real_cmu_aligned_Kuw": { "ann_files": [ osp.join( DATASETS_ROOT, - "BOP_DATASETS/ycbv/image_sets/train_real_cmu.txt", + "BOP_DATASETS/fruitbin/image_sets/train_real_cmu.txt", ) ], - "image_prefixes": [osp.join(DATASETS_ROOT, "BOP_DATASETS/ycbv/train_real")], + "image_prefixes": [osp.join(DATASETS_ROOT, "BOP_DATASETS/fruitbin/train_real")], "align_K_by_change_pose": True, }, - "ycbv_train_synt": { - "ann_files": [osp.join(DATASETS_ROOT, "BOP_DATASETS/ycbv/image_sets/train_synt.txt")], - "image_prefixes": [osp.join(DATASETS_ROOT, "BOP_DATASETS/ycbv/train_synt")], + "fruitbin_train_synt": { + "ann_files": [osp.join(DATASETS_ROOT, "BOP_DATASETS/fruitbin/image_sets/train_synt.txt")], + "image_prefixes": [osp.join(DATASETS_ROOT, "BOP_DATASETS/fruitbin/train_synt")], }, - "ycbv_train_synt_50k": { + "fruitbin_train_synt_50k": { "ann_files": [ osp.join( DATASETS_ROOT, - "BOP_DATASETS/ycbv/image_sets/train_synt_50k.txt", + "BOP_DATASETS/fruitbin/image_sets/train_synt_50k.txt", ) ], - "image_prefixes": [osp.join(DATASETS_ROOT, "BOP_DATASETS/ycbv/train_synt")], + "image_prefixes": [osp.join(DATASETS_ROOT, "BOP_DATASETS/fruitbin/train_synt")], }, - "ycbv_train_synt_30k": { + "fruitbin_train_synt_30k": { "ann_files": [ osp.join( DATASETS_ROOT, - "BOP_DATASETS/ycbv/image_sets/train_synt_30k.txt", + "BOP_DATASETS/fruitbin/image_sets/train_synt_30k.txt", ) ], - "image_prefixes": [osp.join(DATASETS_ROOT, "BOP_DATASETS/ycbv/train_synt")], + "image_prefixes": [osp.join(DATASETS_ROOT, "BOP_DATASETS/fruitbin/train_synt")], }, - "ycbv_train_synt_100": { + "fruitbin_train_synt_100": { "ann_files": [ osp.join( DATASETS_ROOT, - "BOP_DATASETS/ycbv/image_sets/train_synt_100.txt", + "BOP_DATASETS/fruitbin/image_sets/train_synt_100.txt", ) ], - "image_prefixes": [osp.join(DATASETS_ROOT, "BOP_DATASETS/ycbv/train_synt")], + "image_prefixes": [osp.join(DATASETS_ROOT, "BOP_DATASETS/fruitbin/train_synt")], }, - "ycbv_test": { - "ann_files": [osp.join(DATASETS_ROOT, "BOP_DATASETS/ycbv/image_sets/keyframe.txt")], - "image_prefixes": [osp.join(DATASETS_ROOT, "BOP_DATASETS/ycbv/test")], + "fruitbin_test": { + "ann_files": [osp.join(DATASETS_ROOT, "BOP_DATASETS/fruitbin/image_sets/keyframe.txt")], + "image_prefixes": [osp.join(DATASETS_ROOT, "BOP_DATASETS/fruitbin/test")], "with_xyz": False, "filter_invalid": False, }, @@ -497,10 +497,10 @@ for name, update_cfg in update_cfgs.items(): if "_100" in name: num_to_load = 100 used_cfg["num_to_load"] = num_to_load - SPLITS_YCBV[name] = used_cfg + SPLITS_FRUITBIN[name] = used_cfg # single object splits ###################################################### -for obj in ref.ycbv.objects: +for obj in ref.fruitbin.objects: for split in [ "train_real", "train_real_aligned_Kuw", @@ -512,7 +512,7 @@ for obj in ref.ycbv.objects: "train_synt_30k", "test", ]: - name = "ycbv_{}_{}".format(obj, split) + name = "fruitbin_{}_{}".format(obj, split) if split in [ "train_real", "train_real_aligned_Kuw", @@ -552,19 +552,19 @@ for obj in ref.ycbv.objects: } root_name, idx_file = split_idx_file_dict[split] - if name not in SPLITS_YCBV: - SPLITS_YCBV[name] = dict( + if name not in SPLITS_FRUITBIN: + SPLITS_FRUITBIN[name] = dict( name=name, - dataset_root=osp.join(DATASETS_ROOT, "BOP_DATASETS/ycbv/"), - models_root=osp.join(DATASETS_ROOT, "BOP_DATASETS/ycbv/models"), + dataset_root=osp.join(DATASETS_ROOT, "BOP_DATASETS/fruitbin/"), + models_root=osp.join(DATASETS_ROOT, "BOP_DATASETS/fruitbin/models"), objs=[obj], ann_files=[ osp.join( DATASETS_ROOT, - "BOP_DATASETS/ycbv/image_sets/{}".format(idx_file), + "BOP_DATASETS/fruitbin/image_sets/{}".format(idx_file), ) ], - image_prefixes=[osp.join(DATASETS_ROOT, "BOP_DATASETS/ycbv/{}".format(root_name))], + image_prefixes=[osp.join(DATASETS_ROOT, "BOP_DATASETS/fruitbin/{}".format(root_name))], scale_to_meter=0.001, with_masks=True, # (load masks but may not use it) with_depth=True, # (load depth path here, but may not use it) @@ -576,7 +576,7 @@ for obj in ref.ycbv.objects: use_cache=True, num_to_load=-1, filter_invalid=filter_invalid, - ref_key="ycbv", + ref_key="fruitbin", ) @@ -590,32 +590,32 @@ def register_with_name_cfg(name, data_cfg=None): data_cfg can be set in cfg.DATA_CFG.name """ dprint("register dataset: {}".format(name)) - if name in SPLITS_YCBV: - used_cfg = SPLITS_YCBV[name] + if name in SPLITS_FRUITBIN: + used_cfg = SPLITS_FRUITBIN[name] else: assert ( data_cfg is not None - ), f"dataset name {name} is not registered. available datasets: {list(SPLITS_YCBV.keys())}" + ), f"dataset name {name} is not registered. available datasets: {list(SPLITS_FRUITBIN.keys())}" used_cfg = data_cfg - DatasetCatalog.register(name, YCBV_Dataset(used_cfg)) + DatasetCatalog.register(name, FRUITBIN_Dataset(used_cfg)) # something like eval_types MetadataCatalog.get(name).set( - id="ycbv", # NOTE: for pvnet to determine module + id="fruitbin", # NOTE: for pvnet to determine module ref_key=used_cfg["ref_key"], objs=used_cfg["objs"], eval_error_types=["ad", "rete", "proj"], evaluator_type="bop", - **get_ycbv_metadata(obj_names=used_cfg["objs"], ref_key=used_cfg["ref_key"]), + **get_fruitbin_metadata(obj_names=used_cfg["objs"], ref_key=used_cfg["ref_key"]), ) def get_available_datasets(): - return list(SPLITS_YCBV.keys()) + return list(SPLITS_FRUITBIN.keys()) #### tests ############################################### def test_vis(): - # python -m core.datasets.ycbv_d2 ycbv_test + # python -m core.datasets.fruitbin_d2 fruitbin_test dataset_name = sys.argv[1] meta = MetadataCatalog.get(dataset_name) t_start = time.perf_counter() @@ -623,7 +623,7 @@ def test_vis(): with_xyz = False if "test" in dataset_name else True logger.info("Done loading {} samples with {:.3f}s.".format(len(dicts), time.perf_counter() - t_start)) - dirname = "output/ycbv_test-data-vis" + dirname = "output/fruitbin_test-data-vis" os.makedirs(dirname, exist_ok=True) objs = meta.objs for d in dicts: @@ -737,3 +737,4 @@ if __name__ == "__main__": register_with_name_cfg(sys.argv[1]) print("dataset catalog: ", DatasetCatalog.list()) test_vis() + diff --git a/det/yolox/data/datasets/fruitbin_pbr.py b/det/yolox/data/datasets/fruitbin_pbr.py index 9295d4c..93bda78 100644 --- a/det/yolox/data/datasets/fruitbin_pbr.py +++ b/det/yolox/data/datasets/fruitbin_pbr.py @@ -25,7 +25,7 @@ logger = logging.getLogger(__name__) DATASETS_ROOT = osp.normpath(osp.join(PROJ_ROOT, "datasets")) -class YCBV_PBR_Dataset: +class FRUITBIN_PBR_Dataset: def __init__(self, data_cfg): """ Set with_depth and with_masks default to True, @@ -39,7 +39,7 @@ class YCBV_PBR_Dataset: self.dataset_root = data_cfg.get( "dataset_root", - osp.join(DATASETS_ROOT, "BOP_DATASETS/ycbv/train_pbr"), + osp.join(DATASETS_ROOT, "BOP_DATASETS/fruitbin/train_pbr"), ) self.xyz_root = data_cfg.get("xyz_root", osp.join(self.dataset_root, "xyz_crop")) assert osp.exists(self.dataset_root), self.dataset_root @@ -59,14 +59,14 @@ class YCBV_PBR_Dataset: ################################################## # NOTE: careful! Only the selected objects - self.cat_ids = [cat_id for cat_id, obj_name in ref.ycbv.id2obj.items() if obj_name in self.objs] + self.cat_ids = [cat_id for cat_id, obj_name in ref.fruitbin.id2obj.items() if obj_name in self.objs] # map selected objs to [0, num_objs-1] self.cat2label = {v: i for i, v in enumerate(self.cat_ids)} # id_map self.label2cat = {label: cat for cat, label in self.cat2label.items()} self.obj2label = OrderedDict((obj, obj_id) for obj_id, obj in enumerate(self.objs)) ########################################################## - self.scenes = [f"{i:06d}" for i in range(50)] + self.scenes = [f"{i:06d}" for i in range(8)] def __call__(self): """Load light-weight instance annotations of all images into a list of @@ -114,10 +114,10 @@ class YCBV_PBR_Dataset: for str_im_id in tqdm(gt_dict, postfix=f"{scene_id}"): int_im_id = int(str_im_id) - rgb_path = osp.join(scene_root, "rgb/{:06d}.jpg").format(int_im_id) + rgb_path = osp.join(scene_root, "rgb/{:d}.png").format(int_im_id) assert osp.exists(rgb_path), rgb_path - depth_path = osp.join(scene_root, "depth/{:06d}.png".format(int_im_id)) + depth_path = osp.join(scene_root, "depth/{:d}.png".format(int_im_id)) scene_im_id = f"{scene_id}/{int_im_id}" @@ -151,7 +151,7 @@ class YCBV_PBR_Dataset: proj = proj[:2] / proj[2] bbox_visib = gt_info_dict[str_im_id][anno_i]["bbox_visib"] - bbox_obj = gt_info_dict[str_im_id][anno_i]["bbox_obj"] + # bbox_obj = gt_info_dict[str_im_id][anno_i]["bbox_obj"] x1, y1, w, h = bbox_visib if self.filter_invalid: if h <= 1 or w <= 1: @@ -164,9 +164,9 @@ class YCBV_PBR_Dataset: ) mask_visib_file = osp.join( scene_root, - "mask_visib/{:06d}_{:06d}.png".format(int_im_id, anno_i), + "mask_visib/{:d}.png".format(int_im_id), ) - assert osp.exists(mask_file), mask_file + # assert osp.exists(mask_file), mask_file assert osp.exists(mask_visib_file), mask_visib_file # load mask visib TODO: load both mask_visib and mask_full mask_single = mmcv.imread(mask_visib_file, "unchanged") @@ -177,26 +177,23 @@ class YCBV_PBR_Dataset: mask_rle = binary_mask_to_rle(mask_single, compressed=True) # load mask full - mask_full = mmcv.imread(mask_file, "unchanged") - mask_full = mask_full.astype("bool") - mask_full_rle = binary_mask_to_rle(mask_full, compressed=True) + # mask_full = mmcv.imread(mask_file, "unchanged") + # mask_full = mask_full.astype("bool") + # mask_full_rle = binary_mask_to_rle(mask_full, compressed=True) visib_fract = gt_info_dict[str_im_id][anno_i].get("visib_fract", 1.0) - xyz_path = osp.join( - self.xyz_root, - f"{scene_id:06d}/{int_im_id:06d}_{anno_i:06d}-xyz.pkl", - ) + xyz_path = osp.join(self.xyz_root, f"{scene_id:06d}/{int_im_id:06d}_{anno_i:06d}-xyz.pkl") inst = { "category_id": cur_label, # 0-based label - "bbox": bbox_obj, # TODO: load both bbox_obj and bbox_visib + "bbox": bbox_visib, # TODO: load both bbox_obj and bbox_visib "bbox_mode": BoxMode.XYWH_ABS, "pose": pose, "quat": quat, "trans": t, "centroid_2d": proj, # absolute (cx, cy) "segmentation": mask_rle, - "mask_full": mask_full_rle, # TODO: load as mask_full, rle + # "mask_full": mask_full_rle, # TODO: load as mask_full, rle "visib_fract": visib_fract, "xyz_path": xyz_path, } @@ -255,7 +252,7 @@ class YCBV_PBR_Dataset: model = inout.load_ply( osp.join( self.models_root, - f"obj_{ref.ycbv.obj2id[obj_name]:06d}.ply", + f"obj_{ref.fruitbin.obj2id[obj_name]:06d}.ply", ), vertex_scale=self.scale_to_meter, ) @@ -275,7 +272,7 @@ class YCBV_PBR_Dataset: ########### register datasets ############################################################ -def get_ycbv_metadata(obj_names, ref_key): +def get_fruitbin_metadata(obj_names, ref_key): """task specific metadata.""" data_ref = ref.__dict__[ref_key] @@ -296,55 +293,55 @@ def get_ycbv_metadata(obj_names, ref_key): return meta -ycbv_model_root = "BOP_DATASETS/ycbv/models/" +fruitbin_model_root = "BOP_DATASETS/fruitbin/models/" ################################################################################ -SPLITS_YCBV_PBR = dict( - ycbv_train_pbr=dict( - name="ycbv_train_pbr", - objs=ref.ycbv.objects, # selected objects - dataset_root=osp.join(DATASETS_ROOT, "BOP_DATASETS/ycbv/train_pbr"), - models_root=osp.join(DATASETS_ROOT, "BOP_DATASETS/ycbv/models"), - xyz_root=osp.join(DATASETS_ROOT, "BOP_DATASETS/ycbv/train_pbr/xyz_crop"), +SPLITS_FRUITBIN_PBR = dict( + fruitbin_train_pbr=dict( + name="fruitbin_train_pbr", + objs=ref.fruitbin.objects, # selected objects + dataset_root=osp.join(DATASETS_ROOT, "BOP_DATASETS/fruitbin/train_pbr"), + models_root=osp.join(DATASETS_ROOT, "BOP_DATASETS/fruitbin/models"), + xyz_root=osp.join(DATASETS_ROOT, "BOP_DATASETS/fruitbin/train_pbr/xyz_crop"), scale_to_meter=0.001, with_masks=True, # (load masks but may not use it) with_depth=True, # (load depth path here, but may not use it) height=480, width=640, - use_cache=True, + use_cache=False, num_to_load=-1, filter_invalid=True, - ref_key="ycbv", + ref_key="fruitbin", ) ) # single obj splits -for obj in ref.ycbv.objects: +for obj in ref.fruitbin.objects: for split in ["train_pbr"]: - name = "ycbv_{}_{}".format(obj, split) + name = "fruitbin_{}_{}".format(obj, split) if split in ["train_pbr"]: filter_invalid = True elif split in ["test"]: filter_invalid = False else: raise ValueError("{}".format(split)) - if name not in SPLITS_YCBV_PBR: - SPLITS_YCBV_PBR[name] = dict( + if name not in SPLITS_FRUITBIN_PBR: + SPLITS_FRUITBIN_PBR[name] = dict( name=name, objs=[obj], # only this obj - dataset_root=osp.join(DATASETS_ROOT, "BOP_DATASETS/ycbv/train_pbr"), - models_root=osp.join(DATASETS_ROOT, "BOP_DATASETS/ycbv/models"), - xyz_root=osp.join(DATASETS_ROOT, "BOP_DATASETS/ycbv/train_pbr/xyz_crop"), + dataset_root=osp.join(DATASETS_ROOT, "BOP_DATASETS/fruitbin/train_pbr"), + models_root=osp.join(DATASETS_ROOT, "BOP_DATASETS/fruitbin/models"), + xyz_root=osp.join(DATASETS_ROOT, "BOP_DATASETS/fruitbin/train_pbr/xyz_crop"), scale_to_meter=0.001, with_masks=True, # (load masks but may not use it) with_depth=True, # (load depth path here, but may not use it) height=480, width=640, - use_cache=True, + use_cache=False, num_to_load=-1, filter_invalid=filter_invalid, - ref_key="ycbv", + ref_key="fruitbin", ) @@ -358,25 +355,25 @@ def register_with_name_cfg(name, data_cfg=None): data_cfg can be set in cfg.DATA_CFG.name """ dprint("register dataset: {}".format(name)) - if name in SPLITS_YCBV_PBR: - used_cfg = SPLITS_YCBV_PBR[name] + if name in SPLITS_FRUITBIN_PBR: + used_cfg = SPLITS_FRUITBIN_PBR[name] else: assert data_cfg is not None, f"dataset name {name} is not registered" used_cfg = data_cfg - DatasetCatalog.register(name, YCBV_PBR_Dataset(used_cfg)) + DatasetCatalog.register(name, FRUITBIN_PBR_Dataset(used_cfg)) # something like eval_types MetadataCatalog.get(name).set( - id="ycbv", # NOTE: for pvnet to determine module + id="fruitbin", # NOTE: for pvnet to determine module ref_key=used_cfg["ref_key"], objs=used_cfg["objs"], eval_error_types=["ad", "rete", "proj"], evaluator_type="bop", - **get_ycbv_metadata(obj_names=used_cfg["objs"], ref_key=used_cfg["ref_key"]), + **get_fruitbin_metadata(obj_names=used_cfg["objs"], ref_key=used_cfg["ref_key"]), ) def get_available_datasets(): - return list(SPLITS_YCBV_PBR.keys()) + return list(SPLITS_FRUITBIN_PBR.keys()) #### tests ############################################### @@ -474,7 +471,7 @@ if __name__ == "__main__": """Test the dataset loader. Usage: - python -m this_module ycbv_pbr_train + python -m this_module fruitbin_pbr_train """ from lib.vis_utils.image import grid_show from lib.utils.setup_logger import setup_my_logger diff --git a/ref/__init__.py b/ref/__init__.py index 7072f6c..97d2241 100644 --- a/ref/__init__.py +++ b/ref/__init__.py @@ -1 +1 @@ -from . import lmo_full, ycbv, hb, hb_bop19, tudl, tless, icbin, itodd, lm_full +from . import lmo_full, ycbv, hb, hb_bop19, tudl, tless, icbin, itodd, lm_full, fruitbin diff --git a/ref/fruitbin.py b/ref/fruitbin.py index 83502f9..e68ea7c 100644 --- a/ref/fruitbin.py +++ b/ref/fruitbin.py @@ -18,20 +18,21 @@ data_root = osp.join(root_dir, "datasets") bop_root = osp.join(data_root, "BOP_DATASETS/") # ---------------------------------------------------------------- # -# YCBV DATASET +# FRUITBIN DATASET # ---------------------------------------------------------------- # -dataset_root = osp.join(bop_root, "ycbv") +dataset_root = osp.join(bop_root, "fruitbin") -train_real_dir = osp.join(dataset_root, "train_real") -train_render_dir = osp.join(dataset_root, "train_synt") -train_pbr_dir = osp.join(dataset_root, "train_pbr") +train_pbr_dir = dataset_root +#train_real_dir = osp.join(dataset_root, "train_real") +#train_render_dir = osp.join(dataset_root, "train_synt") +#train_pbr_dir = osp.join(dataset_root, "train_pbr") -test_dir = osp.join(dataset_root, "test") +#test_dir = osp.join(dataset_root, "test") -test_scenes = [i for i in range(48, 59 + 1)] -train_real_scenes = [i for i in range(0, 91 + 1) if i not in test_scenes] -train_synt_scenes = [i for i in range(0, 79 + 1)] -train_pbr_scenes = [i for i in range(0, 49 + 1)] +#test_scenes = [i for i in range(48, 59 + 1)] +#train_real_scenes = [i for i in range(0, 91 + 1) if i not in test_scenes] +#train_synt_scenes = [i for i in range(0, 79 + 1)] +#train_pbr_scenes = [i for i in range(0, 49 + 1)] model_dir = osp.join(dataset_root, "models") fine_model_dir = osp.join(dataset_root, "models_fine") @@ -41,43 +42,29 @@ vertex_scale = 0.001 # object info id2obj = { - 1: "002_master_chef_can", # [1.3360, -0.5000, 3.5105] - 2: "003_cracker_box", # [0.5575, 1.7005, 4.8050] - 3: "004_sugar_box", # [-0.9520, 1.4670, 4.3645] - 4: "005_tomato_soup_can", # [-0.0240, -1.5270, 8.4035] - 5: "006_mustard_bottle", # [1.2995, 2.4870, -11.8290] - 6: "007_tuna_fish_can", # [-0.1565, 0.1150, 4.2625] - 7: "008_pudding_box", # [1.1645, -4.2015, 3.1190] - 8: "009_gelatin_box", # [1.4460, -0.5915, 3.6085] - 9: "010_potted_meat_can", # [2.4195, 0.3075, 8.0715] - 10: "011_banana", # [-18.6730, 12.1915, -1.4635] - 11: "019_pitcher_base", # [5.3370, 5.8855, 25.6115] - 12: "021_bleach_cleanser", # [4.9290, -2.4800, -13.2920] - 13: "024_bowl", # [-0.2270, 0.7950, -2.9675] - 14: "025_mug", # [-8.4675, -0.6995, -1.6145] - 15: "035_power_drill", # [9.0710, 20.9360, -2.1190] - 16: "036_wood_block", # [1.4265, -2.5305, 17.1890] - 17: "037_scissors", # [7.0535, -28.1320, 0.0420] - 18: "040_large_marker", # [0.0460, -2.1040, 0.3500] - 19: "051_large_clamp", # [10.5180, -1.9640, -0.4745] - 20: "052_extra_large_clamp", # [-0.3950, -10.4130, 0.1620] - 21: "061_foam_brick", # [-0.0805, 0.0805, -8.2435] + 1: "apple2", + 2: "apricot", + 3: "banana1", + 4: "kiwi1", + 5: "lemon2", + 6: "orange2", + 7: "peach1", + 8: "pear2" } + + objects = list(id2obj.values()) obj_num = len(id2obj) obj2id = {_name: _id for _id, _name in id2obj.items()} -model_paths = [osp.join(model_dir, "obj_{:06d}.ply").format(_id) for _id in id2obj] # TODO: check this +model_paths = [osp.join(model_dir, "obj_{:06d}.ply").format(_id) for _id in id2obj] # TODO: check this texture_paths = [osp.join(model_dir, "obj_{:06d}.png".format(_id)) for _id in id2obj] model_colors = [((i + 1) * 10, (i + 1) * 10, (i + 1) * 10) for i in range(obj_num)] # for renderer # yapf: disable -diameters = np.array([172.063, 269.573, 198.377, 120.543, 196.463, - 89.797, 142.543, 114.053, 129.540, 197.796, - 259.534, 259.566, 161.922, 124.990, 226.170, - 237.299, 203.973, 121.365, 174.746, 217.094, - 102.903]) / 1000.0 +diameters = np.array([52.94260215554199, 55.64182762994041, 152.6314929798569, 72.65268354988685, 72.58068124508065, + 77.41278909198401, 76.63639853229013, 129.85637252518683]) / 1000.0 # yapf: enable # Camera info width = 640 @@ -103,7 +90,7 @@ def get_models_info(): def get_fps_points(): """key is str(obj_id) generated by - core/gdrn_modeling/tools/ycbv/ycbv_1_compute_fps.py.""" + core/gdrn_modeling/tools/fruitbin/fruitbin_1_compute_fps.py.""" fps_points_path = osp.join(model_dir, "fps_points.pkl") assert osp.exists(fps_points_path), fps_points_path fps_dict = mmcv.load(fps_points_path) @@ -112,8 +99,8 @@ def get_fps_points(): def get_keypoints_3d(): """key is str(obj_id) generated by - core/roi_pvnet/tools/ycbv/ycbv_1_compute_keypoints_3d.py.""" + core/roi_pvnet/tools/fruitbin/fruitbin_1_compute_keypoints_3d.py.""" keypoints_3d_path = osp.join(model_dir, "keypoints_3d.pkl") assert osp.exists(keypoints_3d_path), keypoints_3d_path kpts_dict = mmcv.load(keypoints_3d_path) - return kpts_dict + return kpts_dict \ No newline at end of file -- GitLab