Skip to content
Snippets Groups Projects
Commit 7b728c86 authored by liuxingyu's avatar liuxingyu
Browse files

replace lm_d2 with lmo_bop_test

parent 6a65126b
No related branches found
No related tags found
No related merge requests found
Hide whitespace changes
Inline Side-by-side
core/gdrn_modeling/datasets/dataset_factory.py
+ 2
− 2
View file @ 7b728c86
...@@ -8,8 +8,8 @@ import detectron2.utils.comm as comm ...@@ -8,8 +8,8 @@ import detectron2.utils.comm as comm
import ref import ref
from detectron2.data import DatasetCatalog, MetadataCatalog from detectron2.data import DatasetCatalog, MetadataCatalog
from core.gdrn_modeling.datasets import ( from core.gdrn_modeling.datasets import (
lm_dataset_d2,
lm_pbr, lm_pbr,
lmo_bop_test,
ycbv_pbr, ycbv_pbr,
ycbv_d2, ycbv_d2,
ycbv_bop_test, ycbv_bop_test,
...@@ -39,8 +39,8 @@ __all__ = [ ...@@ -39,8 +39,8 @@ __all__ = [
"get_available_datasets", "get_available_datasets",
] ]
_DSET_MOD_NAMES = [ _DSET_MOD_NAMES = [
"lm_dataset_d2",
"lm_pbr", "lm_pbr",
"lmo_bop_test",
"ycbv_pbr", "ycbv_pbr",
"ycbv_d2", "ycbv_d2",
"ycbv_bop_test", "ycbv_bop_test",
......
core/gdrn_modeling/datasets/lm_dataset_d2.py → core/gdrn_modeling/datasets/lmo_bop_test.py
+ 521
− 0
View file @ 7b728c86
...@@ -27,8 +27,8 @@ logger = logging.getLogger(__name__) ...@@ -27,8 +27,8 @@ logger = logging.getLogger(__name__)
DATASETS_ROOT = osp.normpath(osp.join(PROJ_ROOT, "datasets")) DATASETS_ROOT = osp.normpath(osp.join(PROJ_ROOT, "datasets"))
class LM_Dataset(object): class LMO_BOP_TEST_Dataset(object):
"""lm splits.""" """lmo bop test splits."""
def __init__(self, data_cfg): def __init__(self, data_cfg):
""" """
...@@ -41,17 +41,16 @@ class LM_Dataset(object): ...@@ -41,17 +41,16 @@ class LM_Dataset(object):
self.objs = data_cfg["objs"] # selected objects self.objs = data_cfg["objs"] # selected objects
self.ann_files = data_cfg["ann_files"] # idx files with image ids self.dataset_root = data_cfg.get("dataset_root", osp.join(DATASETS_ROOT, "BOP_DATASETS/lmo/test"))
self.image_prefixes = data_cfg["image_prefixes"]
self.xyz_prefixes = data_cfg["xyz_prefixes"]
self.dataset_root = data_cfg["dataset_root"] # BOP_DATASETS/lm/
assert osp.exists(self.dataset_root), self.dataset_root assert osp.exists(self.dataset_root), self.dataset_root
self.models_root = data_cfg["models_root"] # BOP_DATASETS/lm/models
self.ann_file = data_cfg["ann_file"] # json file with scene_id and im_id items
self.models_root = data_cfg["models_root"] # BOP_DATASETS/lmo/models
self.scale_to_meter = data_cfg["scale_to_meter"] # 0.001 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) self.with_masks = data_cfg["with_masks"]
self.with_depth = data_cfg["with_depth"] # True (load depth path here, but may not use it) self.with_depth = data_cfg["with_depth"]
self.height = data_cfg["height"] # 480 self.height = data_cfg["height"] # 480
self.width = data_cfg["width"] # 640 self.width = data_cfg["width"] # 640
...@@ -59,20 +58,18 @@ class LM_Dataset(object): ...@@ -59,20 +58,18 @@ class LM_Dataset(object):
self.cache_dir = data_cfg.get("cache_dir", osp.join(PROJ_ROOT, ".cache")) # .cache self.cache_dir = data_cfg.get("cache_dir", osp.join(PROJ_ROOT, ".cache")) # .cache
self.use_cache = data_cfg.get("use_cache", True) self.use_cache = data_cfg.get("use_cache", True)
self.num_to_load = data_cfg["num_to_load"] # -1 self.num_to_load = data_cfg["num_to_load"] # -1
self.filter_invalid = data_cfg["filter_invalid"] self.filter_invalid = data_cfg.get("filter_invalid", True)
self.filter_scene = data_cfg.get("filter_scene", False)
self.debug_im_id = data_cfg.get("debug_im_id", None)
################################################## ##################################################
# NOTE: careful! Only the selected objects # NOTE: careful! Only the selected objects
self.cat_ids = [cat_id for cat_id, obj_name in ref.lm_full.id2obj.items() if obj_name in self.objs] self.cat_ids = [cat_id for cat_id, obj_name in ref.lmo_full.id2obj.items() if obj_name in self.objs]
# map selected objs to [0, num_objs-1] # map selected objs to [0, num_objs-1]
self.cat2label = {v: i for i, v in enumerate(self.cat_ids)} # id_map 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.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.obj2label = OrderedDict((obj, obj_id) for obj, obj_id in enumerate(self.objs))
########################################################## ##########################################################
def __call__(self): # LM_Dataset def __call__(self):
"""Load light-weight instance annotations of all images into a list of """Load light-weight instance annotations of all images into a list of
dicts in Detectron2 format. dicts in Detectron2 format.
...@@ -92,10 +89,7 @@ class LM_Dataset(object): ...@@ -92,10 +89,7 @@ class LM_Dataset(object):
) )
).encode("utf-8") ).encode("utf-8")
).hexdigest() ).hexdigest()
cache_path = osp.join( cache_path = osp.join(self.cache_dir, "dataset_dicts_{}_{}.pkl".format(self.name, hashed_file_name))
self.cache_dir,
"dataset_dicts_{}_{}.pkl".format(self.name, hashed_file_name),
)
if osp.exists(cache_path) and self.use_cache: if osp.exists(cache_path) and self.use_cache:
logger.info("load cached dataset dicts from {}".format(cache_path)) logger.info("load cached dataset dicts from {}".format(cache_path))
...@@ -107,134 +101,137 @@ class LM_Dataset(object): ...@@ -107,134 +101,137 @@ class LM_Dataset(object):
self.num_instances_without_valid_segmentation = 0 self.num_instances_without_valid_segmentation = 0
self.num_instances_without_valid_box = 0 self.num_instances_without_valid_box = 0
dataset_dicts = [] # ###################################################### dataset_dicts = [] # ######################################################
assert len(self.ann_files) == len(self.image_prefixes), f"{len(self.ann_files)} != {len(self.image_prefixes)}" # it is slow because of loading and converting masks to rle
assert len(self.ann_files) == len(self.xyz_prefixes), f"{len(self.ann_files)} != {len(self.xyz_prefixes)}" targets = mmcv.load(self.ann_file)
unique_im_id = 0
for ann_file, scene_root, xyz_root in zip(tqdm(self.ann_files), self.image_prefixes, self.xyz_prefixes): scene_im_ids = [(item["scene_id"], item["im_id"]) for item in targets]
# linemod each scene is an object scene_im_ids = sorted(list(set(scene_im_ids)))
with open(ann_file, "r") as f_ann:
indices = [line.strip("\r\n") for line in f_ann.readlines()] # string ids # load infos for each scene
gt_dict = mmcv.load(osp.join(scene_root, "scene_gt.json")) gt_dicts = {}
gt_info_dict = mmcv.load(osp.join(scene_root, "scene_gt_info.json")) # bbox_obj, bbox_visib gt_info_dicts = {}
cam_dict = mmcv.load(osp.join(scene_root, "scene_camera.json")) cam_dicts = {}
for im_id in tqdm(indices): for scene_id, im_id in scene_im_ids:
int_im_id = int(im_id) scene_root = osp.join(self.dataset_root, f"{scene_id:06d}")
str_im_id = str(int_im_id) if scene_id not in gt_dicts:
rgb_path = osp.join(scene_root, "rgb/{:06d}.png").format(int_im_id) gt_dicts[scene_id] = mmcv.load(osp.join(scene_root, "scene_gt.json"))
assert osp.exists(rgb_path), rgb_path if scene_id not in gt_info_dicts:
gt_info_dicts[scene_id] = mmcv.load(osp.join(scene_root, "scene_gt_info.json")) # bbox_obj, bbox_visib
depth_path = osp.join(scene_root, "depth/{:06d}.png".format(int_im_id)) if scene_id not in cam_dicts:
cam_dicts[scene_id] = mmcv.load(osp.join(scene_root, "scene_camera.json"))
scene_id = int(rgb_path.split("/")[-3])
scene_im_id = f"{scene_id}/{int_im_id}" for scene_id, int_im_id in tqdm(scene_im_ids):
str_im_id = str(int_im_id)
if self.debug_im_id is not None: scene_root = osp.join(self.dataset_root, f"{scene_id:06d}")
if self.debug_im_id != scene_im_id:
gt_dict = gt_dicts[scene_id]
gt_info_dict = gt_info_dicts[scene_id]
cam_dict = cam_dicts[scene_id]
rgb_path = osp.join(scene_root, "rgb/{:06d}.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))
scene_im_id = f"{scene_id}/{int_im_id}"
K = np.array(cam_dict[str_im_id]["cam_K"], dtype=np.float32).reshape(3, 3)
depth_factor = 1000.0 / cam_dict[str_im_id]["depth_scale"] # 10000
record = {
"dataset_name": self.name,
"file_name": osp.relpath(rgb_path, PROJ_ROOT),
"depth_file": osp.relpath(depth_path, PROJ_ROOT),
"height": self.height,
"width": self.width,
"image_id": int_im_id,
"scene_im_id": scene_im_id, # for evaluation
"cam": K,
"depth_factor": depth_factor,
"img_type": "real", # NOTE: has background
}
insts = []
for anno_i, anno in enumerate(gt_dict[str_im_id]):
obj_id = anno["obj_id"]
if obj_id not in self.cat_ids:
continue
cur_label = self.cat2label[obj_id] # 0-based label
R = np.array(anno["cam_R_m2c"], dtype="float32").reshape(3, 3)
t = np.array(anno["cam_t_m2c"], dtype="float32") / 1000.0
pose = np.hstack([R, t.reshape(3, 1)])
quat = mat2quat(R).astype("float32")
proj = (record["cam"] @ t.T).T
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"]
x1, y1, w, h = bbox_visib
if self.filter_invalid:
if h <= 1 or w <= 1:
self.num_instances_without_valid_box += 1
continue continue
K = np.array(cam_dict[str_im_id]["cam_K"], dtype=np.float32).reshape(3, 3) mask_file = osp.join(
depth_factor = 1000.0 / cam_dict[str_im_id]["depth_scale"] scene_root,
if self.filter_scene: "mask/{:06d}_{:06d}.png".format(int_im_id, anno_i),
if scene_id not in self.cat_ids: )
continue mask_visib_file = osp.join(
record = { scene_root,
"dataset_name": self.name, "mask_visib/{:06d}_{:06d}.png".format(int_im_id, anno_i),
"file_name": osp.relpath(rgb_path, PROJ_ROOT), )
"depth_file": osp.relpath(depth_path, PROJ_ROOT), assert osp.exists(mask_file), mask_file
"height": self.height, assert osp.exists(mask_visib_file), mask_visib_file
"width": self.width, # load mask visib
"image_id": unique_im_id, mask_single = mmcv.imread(mask_visib_file, "unchanged")
"scene_im_id": scene_im_id, # for evaluation mask_single = mask_single.astype("bool")
"cam": K, area = mask_single.sum()
"depth_factor": depth_factor, if area < 3: # filter out too small or nearly invisible instances
"img_type": "real", self.num_instances_without_valid_segmentation += 1
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)
visib_fract = gt_info_dict[str_im_id][anno_i].get("visib_fract", 1.0)
inst = {
"category_id": cur_label, # 0-based label
"bbox": bbox_visib,
"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,
"visib_fract": visib_fract,
"xyz_path": None, # no need for test
} }
unique_im_id += 1
insts = [] model_info = self.models_info[str(obj_id)]
for anno_i, anno in enumerate(gt_dict[str_im_id]): inst["model_info"] = model_info
obj_id = anno["obj_id"] for key in ["bbox3d_and_center"]:
if obj_id not in self.cat_ids: inst[key] = self.models[cur_label][key]
continue insts.append(inst)
cur_label = self.cat2label[obj_id] # 0-based label if len(insts) == 0: # filter im without anno
R = np.array(anno["cam_R_m2c"], dtype="float32").reshape(3, 3) continue
t = np.array(anno["cam_t_m2c"], dtype="float32") / 1000.0 record["annotations"] = insts
pose = np.hstack([R, t.reshape(3, 1)]) dataset_dicts.append(record)
quat = mat2quat(R).astype("float32")
proj = (record["cam"] @ t.T).T
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"]
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(int_im_id, anno_i),
)
mask_visib_file = osp.join(
scene_root,
"mask_visib/{:06d}_{:06d}.png".format(int_im_id, anno_i),
)
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")
mask_single = mask_single.astype("bool")
area = mask_single.sum()
if area < 3: # filter out too small or nearly invisible instances
self.num_instances_without_valid_segmentation += 1
continue
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)
inst = {
"category_id": cur_label, # 0-based label
"bbox": bbox_visib,
"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,
}
if "test" not in self.name.lower():
# if True:
xyz_path = osp.join(xyz_root, f"{int_im_id:06d}_{anno_i:06d}.pkl")
assert osp.exists(xyz_path), xyz_path
inst["xyz_path"] = xyz_path
model_info = self.models_info[str(obj_id)]
inst["model_info"] = model_info
# TODO: using full mask and full xyz
for key in ["bbox3d_and_center"]:
inst[key] = self.models[cur_label][key]
insts.append(inst)
if len(insts) == 0: # filter im without anno
continue
record["annotations"] = insts
dataset_dicts.append(record)
if self.num_instances_without_valid_segmentation > 0: if self.num_instances_without_valid_segmentation > 0:
logger.warning( logger.warning(
"Filtered out {} instances without valid segmentation. " "There are {} instances without valid segmentation. "
"There might be issues in your dataset generation process.".format( "There might be issues in your dataset generation process.".format(
self.num_instances_without_valid_segmentation self.num_instances_without_valid_segmentation
) )
) )
if self.num_instances_without_valid_box > 0: if self.num_instances_without_valid_box > 0:
logger.warning( logger.warning(
"Filtered out {} instances without valid box. " "There are {} instances without valid box. "
"There might be issues in your dataset generation process.".format(self.num_instances_without_valid_box) "There might be issues in your dataset generation process.".format(self.num_instances_without_valid_box)
) )
########################################################################## ##########################################################################
...@@ -268,7 +265,7 @@ class LM_Dataset(object): ...@@ -268,7 +265,7 @@ class LM_Dataset(object):
model = inout.load_ply( model = inout.load_ply(
osp.join( osp.join(
self.models_root, self.models_root,
f"obj_{ref.lm_full.obj2id[obj_name]:06d}.ply", f"obj_{ref.lmo_full.obj2id[obj_name]:06d}.ply",
), ),
vertex_scale=self.scale_to_meter, vertex_scale=self.scale_to_meter,
) )
...@@ -282,6 +279,9 @@ class LM_Dataset(object): ...@@ -282,6 +279,9 @@ class LM_Dataset(object):
mmcv.dump(models, cache_path, protocol=4) mmcv.dump(models, cache_path, protocol=4)
return models return models
def __len__(self):
return self.num_to_load
def image_aspect_ratio(self): def image_aspect_ratio(self):
return self.width / self.height # 4/3 return self.width / self.height # 4/3
...@@ -289,7 +289,7 @@ class LM_Dataset(object): ...@@ -289,7 +289,7 @@ class LM_Dataset(object):
########### register datasets ############################################################ ########### register datasets ############################################################
def get_lm_metadata(obj_names, ref_key): def get_lmo_metadata(obj_names, ref_key):
"""task specific metadata.""" """task specific metadata."""
data_ref = ref.__dict__[ref_key] data_ref = ref.__dict__[ref_key]
...@@ -311,202 +311,15 @@ def get_lm_metadata(obj_names, ref_key): ...@@ -311,202 +311,15 @@ def get_lm_metadata(obj_names, ref_key):
return meta return meta
LM_13_OBJECTS = [ ##########################################################################
"ape",
"benchvise", SPLITS_LMO = dict(
"camera",
"can",
"cat",
"driller",
"duck",
"eggbox",
"glue",
"holepuncher",
"iron",
"lamp",
"phone",
] # no bowl, cup
LM_OCC_OBJECTS = [
"ape",
"can",
"cat",
"driller",
"duck",
"eggbox",
"glue",
"holepuncher",
]
################################################################################
SPLITS_LM = dict(
lm_13_train=dict(
name="lm_13_train",
dataset_root=osp.join(DATASETS_ROOT, "BOP_DATASETS/lm/"),
models_root=osp.join(DATASETS_ROOT, "BOP_DATASETS/lm/models"),
objs=LM_13_OBJECTS, # selected objects
ann_files=[
osp.join(
DATASETS_ROOT,
"BOP_DATASETS/lm/image_set/{}_{}.txt".format(_obj, "train"),
)
for _obj in LM_13_OBJECTS
],
image_prefixes=[
osp.join(
DATASETS_ROOT,
"BOP_DATASETS/lm/test/{:06d}".format(ref.lm_full.obj2id[_obj]),
)
for _obj in LM_13_OBJECTS
],
xyz_prefixes=[
osp.join(
DATASETS_ROOT,
"BOP_DATASETS/lm/test/xyz_crop/{:06d}".format(ref.lm_full.obj2id[_obj]),
)
for _obj in LM_13_OBJECTS
],
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,
num_to_load=-1,
filter_scene=True,
filter_invalid=True,
ref_key="lm_full",
),
lm_13_test=dict(
name="lm_13_test",
dataset_root=osp.join(DATASETS_ROOT, "BOP_DATASETS/lm/"),
models_root=osp.join(DATASETS_ROOT, "BOP_DATASETS/lm/models"),
objs=LM_13_OBJECTS,
ann_files=[
osp.join(
DATASETS_ROOT,
"BOP_DATASETS/lm/image_set/{}_{}.txt".format(_obj, "test"),
)
for _obj in LM_13_OBJECTS
],
# NOTE: scene root
image_prefixes=[
osp.join(DATASETS_ROOT, "BOP_DATASETS/lm/test/{:06d}").format(ref.lm_full.obj2id[_obj])
for _obj in LM_13_OBJECTS
],
xyz_prefixes=[
osp.join(
DATASETS_ROOT,
"BOP_DATASETS/lm/test/xyz_crop/{:06d}".format(ref.lm_full.obj2id[_obj]),
)
for _obj in LM_13_OBJECTS
],
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,
num_to_load=-1,
filter_scene=True,
filter_invalid=False,
ref_key="lm_full",
),
lmo_train=dict(
name="lmo_train",
# use lm real all (8 objects) to train for lmo
dataset_root=osp.join(DATASETS_ROOT, "BOP_DATASETS/lm/"),
models_root=osp.join(DATASETS_ROOT, "BOP_DATASETS/lm/models"),
objs=LM_OCC_OBJECTS, # selected objects
ann_files=[
osp.join(
DATASETS_ROOT,
"BOP_DATASETS/lm/image_set/{}_{}.txt".format(_obj, "all"),
)
for _obj in LM_OCC_OBJECTS
],
image_prefixes=[
osp.join(
DATASETS_ROOT,
"BOP_DATASETS/lm/test/{:06d}".format(ref.lmo_full.obj2id[_obj]),
)
for _obj in LM_OCC_OBJECTS
],
xyz_prefixes=[
osp.join(
DATASETS_ROOT,
"BOP_DATASETS/lm/test/xyz_crop/{:06d}".format(ref.lmo_full.obj2id[_obj]),
)
for _obj in LM_OCC_OBJECTS
],
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,
num_to_load=-1,
filter_scene=True,
filter_invalid=True,
ref_key="lmo_full",
),
lmo_NoBopTest_train=dict(
name="lmo_NoBopTest_train",
dataset_root=osp.join(DATASETS_ROOT, "BOP_DATASETS/lmo/"),
models_root=osp.join(DATASETS_ROOT, "BOP_DATASETS/lmo/models"),
objs=LM_OCC_OBJECTS,
ann_files=[osp.join(DATASETS_ROOT, "BOP_DATASETS/lmo/image_set/lmo_no_bop_test.txt")],
image_prefixes=[osp.join(DATASETS_ROOT, "BOP_DATASETS/lmo/test/{:06d}").format(2)],
xyz_prefixes=[
osp.join(
DATASETS_ROOT,
"BOP_DATASETS/lmo/test/xyz_crop/{:06d}".format(2),
)
],
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,
num_to_load=-1,
filter_scene=False,
filter_invalid=True,
ref_key="lmo_full",
),
lmo_test=dict(
name="lmo_test",
dataset_root=osp.join(DATASETS_ROOT, "BOP_DATASETS/lmo/"),
models_root=osp.join(DATASETS_ROOT, "BOP_DATASETS/lmo/models"),
objs=LM_OCC_OBJECTS,
ann_files=[osp.join(DATASETS_ROOT, "BOP_DATASETS/lmo/image_set/lmo_test.txt")],
# NOTE: scene root
image_prefixes=[osp.join(DATASETS_ROOT, "BOP_DATASETS/lmo/test/{:06d}").format(2)],
xyz_prefixes=[None],
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,
num_to_load=-1,
filter_scene=False,
filter_invalid=False,
ref_key="lmo_full",
),
lmo_bop_test=dict( lmo_bop_test=dict(
name="lmo_bop_test", name="lmo_bop_test",
dataset_root=osp.join(DATASETS_ROOT, "BOP_DATASETS/lmo/"), dataset_root=osp.join(DATASETS_ROOT, "BOP_DATASETS/lmo/test"),
models_root=osp.join(DATASETS_ROOT, "BOP_DATASETS/lmo/models"), models_root=osp.join(DATASETS_ROOT, "BOP_DATASETS/lmo/models"),
objs=LM_OCC_OBJECTS, objs=ref.lmo_full.objects,
ann_files=[osp.join(DATASETS_ROOT, "BOP_DATASETS/lmo/image_set/lmo_bop_test.txt")], ann_file=osp.join(DATASETS_ROOT, "BOP_DATASETS/lmo/test_targets_bop19.json"),
# NOTE: scene root
image_prefixes=[osp.join(DATASETS_ROOT, "BOP_DATASETS/lmo/test/{:06d}").format(2)],
xyz_prefixes=[None],
scale_to_meter=0.001, scale_to_meter=0.001,
with_masks=True, # (load masks but may not use it) with_masks=True, # (load masks but may not use it)
with_depth=True, # (load depth path here, but may not use it) with_depth=True, # (load depth path here, but may not use it)
...@@ -515,160 +328,31 @@ SPLITS_LM = dict( ...@@ -515,160 +328,31 @@ SPLITS_LM = dict(
cache_dir=osp.join(PROJ_ROOT, ".cache"), cache_dir=osp.join(PROJ_ROOT, ".cache"),
use_cache=True, use_cache=True,
num_to_load=-1, num_to_load=-1,
filter_scene=False,
filter_invalid=False, filter_invalid=False,
ref_key="lmo_full", ref_key="lmo_full",
), ),
) )
# single obj splits for lm real # single obj splits for lmo bop test
for obj in ref.lm_full.objects: for obj in ref.lmo_full.objects:
for split in ["train", "test", "all"]: for split in [
name = "lm_real_{}_{}".format(obj, split) "bop_test",
]:
name = "lmo_{}_{}".format(obj, split)
ann_files = [ ann_files = [
osp.join( osp.join(
DATASETS_ROOT, DATASETS_ROOT,
"BOP_DATASETS/lm/image_set/{}_{}.txt".format(obj, split), "BOP_DATASETS/lmo/image_set/{}_{}.txt".format(obj, split),
) )
] ]
if split in ["train", "all"]: # all is used to train lmo if name not in SPLITS_LMO:
filter_invalid = True SPLITS_LMO[name] = dict(
elif split in ["test"]:
filter_invalid = False
else:
raise ValueError("{}".format(split))
if name not in SPLITS_LM:
SPLITS_LM[name] = dict(
name=name,
dataset_root=osp.join(DATASETS_ROOT, "BOP_DATASETS/lm/"),
models_root=osp.join(DATASETS_ROOT, "BOP_DATASETS/lm/models"),
objs=[obj], # only this obj
ann_files=ann_files,
image_prefixes=[osp.join(DATASETS_ROOT, "BOP_DATASETS/lm/test/{:06d}").format(ref.lm_full.obj2id[obj])],
xyz_prefixes=[
osp.join(
DATASETS_ROOT,
"BOP_DATASETS/lm/test/xyz_crop/{:06d}".format(ref.lm_full.obj2id[obj]),
)
],
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,
num_to_load=-1,
filter_invalid=filter_invalid,
filter_scene=True,
ref_key="lm_full",
)
# single obj splits for lmo_NoBopTest_train
for obj in ref.lmo_full.objects:
for split in ["train"]:
name = "lmo_NoBopTest_{}_{}".format(obj, split)
if split in ["train"]:
filter_invalid = True
else:
raise ValueError("{}".format(split))
if name not in SPLITS_LM:
SPLITS_LM[name] = dict(
name=name,
dataset_root=osp.join(DATASETS_ROOT, "BOP_DATASETS/lmo/"),
models_root=osp.join(DATASETS_ROOT, "BOP_DATASETS/lmo/models"),
objs=[obj],
ann_files=[
osp.join(
DATASETS_ROOT,
"BOP_DATASETS/lmo/image_set/lmo_no_bop_test.txt",
)
],
# NOTE: scene root
image_prefixes=[osp.join(DATASETS_ROOT, "BOP_DATASETS/lmo/test/{:06d}").format(2)],
xyz_prefixes=[
osp.join(
DATASETS_ROOT,
"BOP_DATASETS/lmo/test/xyz_crop/{:06d}".format(2),
)
],
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,
num_to_load=-1,
filter_scene=False,
filter_invalid=filter_invalid,
ref_key="lmo_full",
)
# single obj splits for lmo_test
for obj in ref.lmo_full.objects:
for split in ["test"]:
name = "lmo_{}_{}".format(obj, split)
if split in ["train", "all"]: # all is used to train lmo
filter_invalid = True
elif split in ["test"]:
filter_invalid = False
else:
raise ValueError("{}".format(split))
if name not in SPLITS_LM:
SPLITS_LM[name] = dict(
name=name, name=name,
dataset_root=osp.join(DATASETS_ROOT, "BOP_DATASETS/lmo/"), dataset_root=osp.join(DATASETS_ROOT, "BOP_DATASETS/lmo/"),
models_root=osp.join(DATASETS_ROOT, "BOP_DATASETS/lmo/models"), models_root=osp.join(DATASETS_ROOT, "BOP_DATASETS/lmo/models"),
objs=[obj], objs=[obj], # only this obj
ann_files=[
osp.join(
DATASETS_ROOT,
"BOP_DATASETS/lmo/image_set/lmo_test.txt",
)
],
# NOTE: scene root
image_prefixes=[osp.join(DATASETS_ROOT, "BOP_DATASETS/lmo/test/{:06d}").format(2)],
xyz_prefixes=[None],
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,
num_to_load=-1,
filter_scene=False,
filter_invalid=False,
ref_key="lmo_full",
)
# single obj splits for lmo_bop_test
for obj in ref.lmo_full.objects:
for split in ["test"]:
name = "lmo_{}_bop_{}".format(obj, split)
if split in ["train", "all"]: # all is used to train lmo
filter_invalid = True
elif split in ["test"]:
filter_invalid = False
else:
raise ValueError("{}".format(split))
if name not in SPLITS_LM:
SPLITS_LM[name] = dict(
name=name,
dataset_root=osp.join(DATASETS_ROOT, "BOP_DATASETS/lmo/"),
models_root=osp.join(DATASETS_ROOT, "BOP_DATASETS/lmo/models"),
objs=[obj],
ann_files=[
osp.join(
DATASETS_ROOT,
"BOP_DATASETS/lmo/image_set/lmo_bop_test.txt",
)
],
# NOTE: scene root
image_prefixes=[osp.join(DATASETS_ROOT, "BOP_DATASETS/lmo/test/{:06d}").format(2)],
xyz_prefixes=[None],
scale_to_meter=0.001, scale_to_meter=0.001,
ann_file=osp.join(DATASETS_ROOT, "BOP_DATASETS/lmo/test_targets_bop19.json"),
with_masks=True, # (load masks but may not use it) with_masks=True, # (load masks but may not use it)
with_depth=True, # (load depth path here, but may not use it) with_depth=True, # (load depth path here, but may not use it)
height=480, height=480,
...@@ -676,60 +360,10 @@ for obj in ref.lmo_full.objects: ...@@ -676,60 +360,10 @@ for obj in ref.lmo_full.objects:
cache_dir=osp.join(PROJ_ROOT, ".cache"), cache_dir=osp.join(PROJ_ROOT, ".cache"),
use_cache=True, use_cache=True,
num_to_load=-1, num_to_load=-1,
filter_scene=False,
filter_invalid=False, filter_invalid=False,
ref_key="lmo_full", ref_key="lmo_full",
) )
# ================ add single image dataset for debug =======================================
debug_im_ids = {
"train": {obj: [] for obj in ref.lm_full.objects},
"test": {obj: [] for obj in ref.lm_full.objects},
}
for obj in ref.lm_full.objects:
for split in ["train", "test"]:
cur_ann_file = osp.join(DATASETS_ROOT, f"BOP_DATASETS/lm/image_set/{obj}_{split}.txt")
ann_files = [cur_ann_file]
im_ids = []
with open(cur_ann_file, "r") as f:
for line in f:
# scene_id(obj_id)/im_id
im_ids.append("{}/{}".format(ref.lm_full.obj2id[obj], int(line.strip("\r\n"))))
debug_im_ids[split][obj] = im_ids
for debug_im_id in debug_im_ids[split][obj]:
name = "lm_single_{}{}_{}".format(obj, debug_im_id.split("/")[1], split)
if name not in SPLITS_LM:
SPLITS_LM[name] = dict(
name=name,
dataset_root=osp.join(DATASETS_ROOT, "BOP_DATASETS/lm/"),
models_root=osp.join(DATASETS_ROOT, "BOP_DATASETS/lm/models"),
objs=[obj], # only this obj
ann_files=ann_files,
image_prefixes=[
osp.join(DATASETS_ROOT, "BOP_DATASETS/lm/test/{:06d}").format(ref.lm_full.obj2id[obj])
],
xyz_prefixes=[
osp.join(
DATASETS_ROOT,
"BOP_DATASETS/lm/test/xyz_crop/{:06d}".format(ref.lm_full.obj2id[obj]),
)
],
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,
num_to_load=-1,
filter_invalid=False,
filter_scene=True,
ref_key="lm_full",
debug_im_id=debug_im_id, # NOTE: debug im id
)
def register_with_name_cfg(name, data_cfg=None): def register_with_name_cfg(name, data_cfg=None):
"""Assume pre-defined datasets live in `./datasets`. """Assume pre-defined datasets live in `./datasets`.
...@@ -741,25 +375,25 @@ def register_with_name_cfg(name, data_cfg=None): ...@@ -741,25 +375,25 @@ def register_with_name_cfg(name, data_cfg=None):
data_cfg can be set in cfg.DATA_CFG.name data_cfg can be set in cfg.DATA_CFG.name
""" """
dprint("register dataset: {}".format(name)) dprint("register dataset: {}".format(name))
if name in SPLITS_LM: if name in SPLITS_LMO:
used_cfg = SPLITS_LM[name] used_cfg = SPLITS_LMO[name]
else: else:
assert data_cfg is not None, f"dataset name {name} is not registered" assert data_cfg is not None, f"dataset name {name} is not registered"
used_cfg = data_cfg used_cfg = data_cfg
DatasetCatalog.register(name, LM_Dataset(used_cfg)) DatasetCatalog.register(name, LMO_BOP_TEST_Dataset(used_cfg))
# something like eval_types # something like eval_types
MetadataCatalog.get(name).set( MetadataCatalog.get(name).set(
id="linemod", # NOTE: for pvnet to determine module id="lmo", # NOTE: for pvnet to determine module
ref_key=used_cfg["ref_key"], ref_key=used_cfg["ref_key"],
objs=used_cfg["objs"], objs=used_cfg["objs"],
eval_error_types=["ad", "rete", "proj"], eval_error_types=["ad", "rete", "proj"],
evaluator_type="bop", evaluator_type="bop",
**get_lm_metadata(obj_names=used_cfg["objs"], ref_key=used_cfg["ref_key"]), **get_lmo_metadata(obj_names=used_cfg["objs"], ref_key=used_cfg["ref_key"]),
) )
def get_available_datasets(): def get_available_datasets():
return list(SPLITS_LM.keys()) return list(SPLITS_LMO.keys())
#### tests ############################################### #### tests ###############################################
......
det/yolox/data/datasets/dataset_factory.py
+ 2
− 2
View file @ 7b728c86
...@@ -3,8 +3,8 @@ import os.path as osp ...@@ -3,8 +3,8 @@ import os.path as osp
import mmcv import mmcv
from detectron2.data import DatasetCatalog from detectron2.data import DatasetCatalog
from . import ( from . import (
lm_dataset_d2,
lm_pbr, lm_pbr,
lmo_bop_test,
ycbv_pbr, ycbv_pbr,
ycbv_d2, ycbv_d2,
ycbv_bop_test, ycbv_bop_test,
...@@ -34,8 +34,8 @@ __all__ = [ ...@@ -34,8 +34,8 @@ __all__ = [
"get_available_datasets", "get_available_datasets",
] ]
_DSET_MOD_NAMES = [ _DSET_MOD_NAMES = [
"lm_dataset_d2",
"lm_pbr", "lm_pbr",
"lmo_bop_test",
"ycbv_pbr", "ycbv_pbr",
"ycbv_d2", "ycbv_d2",
"ycbv_bop_test", "ycbv_bop_test",
......
det/yolox/data/datasets/lm_dataset_d2.py → det/yolox/data/datasets/lmo_bop_test.py
+ 521
− 0
View file @ 7b728c86
...@@ -27,8 +27,8 @@ logger = logging.getLogger(__name__) ...@@ -27,8 +27,8 @@ logger = logging.getLogger(__name__)
DATASETS_ROOT = osp.normpath(osp.join(PROJ_ROOT, "datasets")) DATASETS_ROOT = osp.normpath(osp.join(PROJ_ROOT, "datasets"))
class LM_Dataset(object): class LMO_BOP_TEST_Dataset(object):
"""lm splits.""" """lmo bop test splits."""
def __init__(self, data_cfg): def __init__(self, data_cfg):
""" """
...@@ -41,17 +41,16 @@ class LM_Dataset(object): ...@@ -41,17 +41,16 @@ class LM_Dataset(object):
self.objs = data_cfg["objs"] # selected objects self.objs = data_cfg["objs"] # selected objects
self.ann_files = data_cfg["ann_files"] # idx files with image ids self.dataset_root = data_cfg.get("dataset_root", osp.join(DATASETS_ROOT, "BOP_DATASETS/lmo/test"))
self.image_prefixes = data_cfg["image_prefixes"]
self.xyz_prefixes = data_cfg["xyz_prefixes"]
self.dataset_root = data_cfg["dataset_root"] # BOP_DATASETS/lm/
assert osp.exists(self.dataset_root), self.dataset_root assert osp.exists(self.dataset_root), self.dataset_root
self.models_root = data_cfg["models_root"] # BOP_DATASETS/lm/models
self.ann_file = data_cfg["ann_file"] # json file with scene_id and im_id items
self.models_root = data_cfg["models_root"] # BOP_DATASETS/lmo/models
self.scale_to_meter = data_cfg["scale_to_meter"] # 0.001 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) self.with_masks = data_cfg["with_masks"]
self.with_depth = data_cfg["with_depth"] # True (load depth path here, but may not use it) self.with_depth = data_cfg["with_depth"]
self.height = data_cfg["height"] # 480 self.height = data_cfg["height"] # 480
self.width = data_cfg["width"] # 640 self.width = data_cfg["width"] # 640
...@@ -59,20 +58,18 @@ class LM_Dataset(object): ...@@ -59,20 +58,18 @@ class LM_Dataset(object):
self.cache_dir = data_cfg.get("cache_dir", osp.join(PROJ_ROOT, ".cache")) # .cache self.cache_dir = data_cfg.get("cache_dir", osp.join(PROJ_ROOT, ".cache")) # .cache
self.use_cache = data_cfg.get("use_cache", True) self.use_cache = data_cfg.get("use_cache", True)
self.num_to_load = data_cfg["num_to_load"] # -1 self.num_to_load = data_cfg["num_to_load"] # -1
self.filter_invalid = data_cfg["filter_invalid"] self.filter_invalid = data_cfg.get("filter_invalid", True)
self.filter_scene = data_cfg.get("filter_scene", False)
self.debug_im_id = data_cfg.get("debug_im_id", None)
################################################## ##################################################
# NOTE: careful! Only the selected objects # NOTE: careful! Only the selected objects
self.cat_ids = [cat_id for cat_id, obj_name in ref.lm_full.id2obj.items() if obj_name in self.objs] self.cat_ids = [cat_id for cat_id, obj_name in ref.lmo_full.id2obj.items() if obj_name in self.objs]
# map selected objs to [0, num_objs-1] # map selected objs to [0, num_objs-1]
self.cat2label = {v: i for i, v in enumerate(self.cat_ids)} # id_map 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.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.obj2label = OrderedDict((obj, obj_id) for obj, obj_id in enumerate(self.objs))
########################################################## ##########################################################
def __call__(self): # LM_Dataset def __call__(self):
"""Load light-weight instance annotations of all images into a list of """Load light-weight instance annotations of all images into a list of
dicts in Detectron2 format. dicts in Detectron2 format.
...@@ -92,10 +89,7 @@ class LM_Dataset(object): ...@@ -92,10 +89,7 @@ class LM_Dataset(object):
) )
).encode("utf-8") ).encode("utf-8")
).hexdigest() ).hexdigest()
cache_path = osp.join( cache_path = osp.join(self.cache_dir, "dataset_dicts_{}_{}.pkl".format(self.name, hashed_file_name))
self.cache_dir,
"dataset_dicts_{}_{}.pkl".format(self.name, hashed_file_name),
)
if osp.exists(cache_path) and self.use_cache: if osp.exists(cache_path) and self.use_cache:
logger.info("load cached dataset dicts from {}".format(cache_path)) logger.info("load cached dataset dicts from {}".format(cache_path))
...@@ -107,133 +101,137 @@ class LM_Dataset(object): ...@@ -107,133 +101,137 @@ class LM_Dataset(object):
self.num_instances_without_valid_segmentation = 0 self.num_instances_without_valid_segmentation = 0
self.num_instances_without_valid_box = 0 self.num_instances_without_valid_box = 0
dataset_dicts = [] # ###################################################### dataset_dicts = [] # ######################################################
assert len(self.ann_files) == len(self.image_prefixes), f"{len(self.ann_files)} != {len(self.image_prefixes)}" # it is slow because of loading and converting masks to rle
assert len(self.ann_files) == len(self.xyz_prefixes), f"{len(self.ann_files)} != {len(self.xyz_prefixes)}" targets = mmcv.load(self.ann_file)
unique_im_id = 0
for ann_file, scene_root, xyz_root in zip(tqdm(self.ann_files), self.image_prefixes, self.xyz_prefixes): scene_im_ids = [(item["scene_id"], item["im_id"]) for item in targets]
# linemod each scene is an object scene_im_ids = sorted(list(set(scene_im_ids)))
with open(ann_file, "r") as f_ann:
indices = [line.strip("\r\n") for line in f_ann.readlines()] # string ids # load infos for each scene
gt_dict = mmcv.load(osp.join(scene_root, "scene_gt.json")) gt_dicts = {}
gt_info_dict = mmcv.load(osp.join(scene_root, "scene_gt_info.json")) # bbox_obj, bbox_visib gt_info_dicts = {}
cam_dict = mmcv.load(osp.join(scene_root, "scene_camera.json")) cam_dicts = {}
for im_id in tqdm(indices): for scene_id, im_id in scene_im_ids:
int_im_id = int(im_id) scene_root = osp.join(self.dataset_root, f"{scene_id:06d}")
str_im_id = str(int_im_id) if scene_id not in gt_dicts:
rgb_path = osp.join(scene_root, "rgb/{:06d}.png").format(int_im_id) gt_dicts[scene_id] = mmcv.load(osp.join(scene_root, "scene_gt.json"))
assert osp.exists(rgb_path), rgb_path if scene_id not in gt_info_dicts:
gt_info_dicts[scene_id] = mmcv.load(osp.join(scene_root, "scene_gt_info.json")) # bbox_obj, bbox_visib
depth_path = osp.join(scene_root, "depth/{:06d}.png".format(int_im_id)) if scene_id not in cam_dicts:
cam_dicts[scene_id] = mmcv.load(osp.join(scene_root, "scene_camera.json"))
scene_id = int(rgb_path.split("/")[-3])
scene_im_id = f"{scene_id}/{int_im_id}" for scene_id, int_im_id in tqdm(scene_im_ids):
str_im_id = str(int_im_id)
if self.debug_im_id is not None: scene_root = osp.join(self.dataset_root, f"{scene_id:06d}")
if self.debug_im_id != scene_im_id:
gt_dict = gt_dicts[scene_id]
gt_info_dict = gt_info_dicts[scene_id]
cam_dict = cam_dicts[scene_id]
rgb_path = osp.join(scene_root, "rgb/{:06d}.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))
scene_im_id = f"{scene_id}/{int_im_id}"
K = np.array(cam_dict[str_im_id]["cam_K"], dtype=np.float32).reshape(3, 3)
depth_factor = 1000.0 / cam_dict[str_im_id]["depth_scale"] # 10000
record = {
"dataset_name": self.name,
"file_name": osp.relpath(rgb_path, PROJ_ROOT),
"depth_file": osp.relpath(depth_path, PROJ_ROOT),
"height": self.height,
"width": self.width,
"image_id": int_im_id,
"scene_im_id": scene_im_id, # for evaluation
"cam": K,
"depth_factor": depth_factor,
"img_type": "real", # NOTE: has background
}
insts = []
for anno_i, anno in enumerate(gt_dict[str_im_id]):
obj_id = anno["obj_id"]
if obj_id not in self.cat_ids:
continue
cur_label = self.cat2label[obj_id] # 0-based label
R = np.array(anno["cam_R_m2c"], dtype="float32").reshape(3, 3)
t = np.array(anno["cam_t_m2c"], dtype="float32") / 1000.0
pose = np.hstack([R, t.reshape(3, 1)])
quat = mat2quat(R).astype("float32")
proj = (record["cam"] @ t.T).T
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"]
x1, y1, w, h = bbox_visib
if self.filter_invalid:
if h <= 1 or w <= 1:
self.num_instances_without_valid_box += 1
continue continue
K = np.array(cam_dict[str_im_id]["cam_K"], dtype=np.float32).reshape(3, 3) mask_file = osp.join(
depth_factor = 1000.0 / cam_dict[str_im_id]["depth_scale"] scene_root,
if self.filter_scene: "mask/{:06d}_{:06d}.png".format(int_im_id, anno_i),
if scene_id not in self.cat_ids: )
continue mask_visib_file = osp.join(
record = { scene_root,
"dataset_name": self.name, "mask_visib/{:06d}_{:06d}.png".format(int_im_id, anno_i),
"file_name": osp.relpath(rgb_path, PROJ_ROOT), )
"depth_file": osp.relpath(depth_path, PROJ_ROOT), assert osp.exists(mask_file), mask_file
"height": self.height, assert osp.exists(mask_visib_file), mask_visib_file
"width": self.width, # load mask visib
"image_id": unique_im_id, mask_single = mmcv.imread(mask_visib_file, "unchanged")
"scene_im_id": scene_im_id, # for evaluation mask_single = mask_single.astype("bool")
"cam": K, area = mask_single.sum()
"depth_factor": depth_factor, if area < 3: # filter out too small or nearly invisible instances
"img_type": "real", self.num_instances_without_valid_segmentation += 1
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)
visib_fract = gt_info_dict[str_im_id][anno_i].get("visib_fract", 1.0)
inst = {
"category_id": cur_label, # 0-based label
"bbox": bbox_visib,
"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,
"visib_fract": visib_fract,
"xyz_path": None, # no need for test
} }
unique_im_id += 1
insts = [] model_info = self.models_info[str(obj_id)]
for anno_i, anno in enumerate(gt_dict[str_im_id]): inst["model_info"] = model_info
obj_id = anno["obj_id"] for key in ["bbox3d_and_center"]:
if obj_id not in self.cat_ids: inst[key] = self.models[cur_label][key]
continue insts.append(inst)
cur_label = self.cat2label[obj_id] # 0-based label if len(insts) == 0: # filter im without anno
R = np.array(anno["cam_R_m2c"], dtype="float32").reshape(3, 3) continue
t = np.array(anno["cam_t_m2c"], dtype="float32") / 1000.0 record["annotations"] = insts
pose = np.hstack([R, t.reshape(3, 1)]) dataset_dicts.append(record)
quat = mat2quat(R).astype("float32")
proj = (record["cam"] @ t.T).T
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"]
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(int_im_id, anno_i),
)
mask_visib_file = osp.join(
scene_root,
"mask_visib/{:06d}_{:06d}.png".format(int_im_id, anno_i),
)
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")
mask_single = mask_single.astype("bool")
area = mask_single.sum()
if area < 3: # filter out too small or nearly invisible instances
self.num_instances_without_valid_segmentation += 1
continue
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)
inst = {
"category_id": cur_label, # 0-based label
"bbox": bbox_obj, # 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,
}
if "test" not in self.name.lower():
# if True:
xyz_path = osp.join(xyz_root, f"{int_im_id:06d}_{anno_i:06d}.pkl")
assert osp.exists(xyz_path), xyz_path
inst["xyz_path"] = xyz_path
model_info = self.models_info[str(obj_id)]
inst["model_info"] = model_info
# TODO: using full mask and full xyz
for key in ["bbox3d_and_center"]:
inst[key] = self.models[cur_label][key]
insts.append(inst)
if len(insts) == 0: # filter im without anno
continue
record["annotations"] = insts
dataset_dicts.append(record)
if self.num_instances_without_valid_segmentation > 0: if self.num_instances_without_valid_segmentation > 0:
logger.warning( logger.warning(
"Filtered out {} instances without valid segmentation. " "There are {} instances without valid segmentation. "
"There might be issues in your dataset generation process.".format( "There might be issues in your dataset generation process.".format(
self.num_instances_without_valid_segmentation self.num_instances_without_valid_segmentation
) )
) )
if self.num_instances_without_valid_box > 0: if self.num_instances_without_valid_box > 0:
logger.warning( logger.warning(
"Filtered out {} instances without valid box. " "There are {} instances without valid box. "
"There might be issues in your dataset generation process.".format(self.num_instances_without_valid_box) "There might be issues in your dataset generation process.".format(self.num_instances_without_valid_box)
) )
########################################################################## ##########################################################################
...@@ -267,7 +265,7 @@ class LM_Dataset(object): ...@@ -267,7 +265,7 @@ class LM_Dataset(object):
model = inout.load_ply( model = inout.load_ply(
osp.join( osp.join(
self.models_root, self.models_root,
f"obj_{ref.lm_full.obj2id[obj_name]:06d}.ply", f"obj_{ref.lmo_full.obj2id[obj_name]:06d}.ply",
), ),
vertex_scale=self.scale_to_meter, vertex_scale=self.scale_to_meter,
) )
...@@ -281,6 +279,9 @@ class LM_Dataset(object): ...@@ -281,6 +279,9 @@ class LM_Dataset(object):
mmcv.dump(models, cache_path, protocol=4) mmcv.dump(models, cache_path, protocol=4)
return models return models
def __len__(self):
return self.num_to_load
def image_aspect_ratio(self): def image_aspect_ratio(self):
return self.width / self.height # 4/3 return self.width / self.height # 4/3
...@@ -288,7 +289,7 @@ class LM_Dataset(object): ...@@ -288,7 +289,7 @@ class LM_Dataset(object):
########### register datasets ############################################################ ########### register datasets ############################################################
def get_lm_metadata(obj_names, ref_key): def get_lmo_metadata(obj_names, ref_key):
"""task specific metadata.""" """task specific metadata."""
data_ref = ref.__dict__[ref_key] data_ref = ref.__dict__[ref_key]
...@@ -310,202 +311,15 @@ def get_lm_metadata(obj_names, ref_key): ...@@ -310,202 +311,15 @@ def get_lm_metadata(obj_names, ref_key):
return meta return meta
LM_13_OBJECTS = [ ##########################################################################
"ape",
"benchvise", SPLITS_LMO = dict(
"camera",
"can",
"cat",
"driller",
"duck",
"eggbox",
"glue",
"holepuncher",
"iron",
"lamp",
"phone",
] # no bowl, cup
LM_OCC_OBJECTS = [
"ape",
"can",
"cat",
"driller",
"duck",
"eggbox",
"glue",
"holepuncher",
]
################################################################################
SPLITS_LM = dict(
lm_13_train=dict(
name="lm_13_train",
dataset_root=osp.join(DATASETS_ROOT, "BOP_DATASETS/lm/"),
models_root=osp.join(DATASETS_ROOT, "BOP_DATASETS/lm/models"),
objs=LM_13_OBJECTS, # selected objects
ann_files=[
osp.join(
DATASETS_ROOT,
"BOP_DATASETS/lm/image_set/{}_{}.txt".format(_obj, "train"),
)
for _obj in LM_13_OBJECTS
],
image_prefixes=[
osp.join(
DATASETS_ROOT,
"BOP_DATASETS/lm/test/{:06d}".format(ref.lm_full.obj2id[_obj]),
)
for _obj in LM_13_OBJECTS
],
xyz_prefixes=[
osp.join(
DATASETS_ROOT,
"BOP_DATASETS/lm/test/xyz_crop/{:06d}".format(ref.lm_full.obj2id[_obj]),
)
for _obj in LM_13_OBJECTS
],
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,
num_to_load=-1,
filter_scene=True,
filter_invalid=True,
ref_key="lm_full",
),
lm_13_test=dict(
name="lm_13_test",
dataset_root=osp.join(DATASETS_ROOT, "BOP_DATASETS/lm/"),
models_root=osp.join(DATASETS_ROOT, "BOP_DATASETS/lm/models"),
objs=LM_13_OBJECTS,
ann_files=[
osp.join(
DATASETS_ROOT,
"BOP_DATASETS/lm/image_set/{}_{}.txt".format(_obj, "test"),
)
for _obj in LM_13_OBJECTS
],
# NOTE: scene root
image_prefixes=[
osp.join(DATASETS_ROOT, "BOP_DATASETS/lm/test/{:06d}").format(ref.lm_full.obj2id[_obj])
for _obj in LM_13_OBJECTS
],
xyz_prefixes=[
osp.join(
DATASETS_ROOT,
"BOP_DATASETS/lm/test/xyz_crop/{:06d}".format(ref.lm_full.obj2id[_obj]),
)
for _obj in LM_13_OBJECTS
],
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,
num_to_load=-1,
filter_scene=True,
filter_invalid=False,
ref_key="lm_full",
),
lmo_train=dict(
name="lmo_train",
# use lm real all (8 objects) to train for lmo
dataset_root=osp.join(DATASETS_ROOT, "BOP_DATASETS/lm/"),
models_root=osp.join(DATASETS_ROOT, "BOP_DATASETS/lm/models"),
objs=LM_OCC_OBJECTS, # selected objects
ann_files=[
osp.join(
DATASETS_ROOT,
"BOP_DATASETS/lm/image_set/{}_{}.txt".format(_obj, "all"),
)
for _obj in LM_OCC_OBJECTS
],
image_prefixes=[
osp.join(
DATASETS_ROOT,
"BOP_DATASETS/lm/test/{:06d}".format(ref.lmo_full.obj2id[_obj]),
)
for _obj in LM_OCC_OBJECTS
],
xyz_prefixes=[
osp.join(
DATASETS_ROOT,
"BOP_DATASETS/lm/test/xyz_crop/{:06d}".format(ref.lmo_full.obj2id[_obj]),
)
for _obj in LM_OCC_OBJECTS
],
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,
num_to_load=-1,
filter_scene=True,
filter_invalid=True,
ref_key="lmo_full",
),
lmo_NoBopTest_train=dict(
name="lmo_NoBopTest_train",
dataset_root=osp.join(DATASETS_ROOT, "BOP_DATASETS/lmo/"),
models_root=osp.join(DATASETS_ROOT, "BOP_DATASETS/lmo/models"),
objs=LM_OCC_OBJECTS,
ann_files=[osp.join(DATASETS_ROOT, "BOP_DATASETS/lmo/image_set/lmo_no_bop_test.txt")],
image_prefixes=[osp.join(DATASETS_ROOT, "BOP_DATASETS/lmo/test/{:06d}").format(2)],
xyz_prefixes=[
osp.join(
DATASETS_ROOT,
"BOP_DATASETS/lmo/test/xyz_crop/{:06d}".format(2),
)
],
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,
num_to_load=-1,
filter_scene=False,
filter_invalid=True,
ref_key="lmo_full",
),
lmo_test=dict(
name="lmo_test",
dataset_root=osp.join(DATASETS_ROOT, "BOP_DATASETS/lmo/"),
models_root=osp.join(DATASETS_ROOT, "BOP_DATASETS/lmo/models"),
objs=LM_OCC_OBJECTS,
ann_files=[osp.join(DATASETS_ROOT, "BOP_DATASETS/lmo/image_set/lmo_test.txt")],
# NOTE: scene root
image_prefixes=[osp.join(DATASETS_ROOT, "BOP_DATASETS/lmo/test/{:06d}").format(2)],
xyz_prefixes=[None],
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,
num_to_load=-1,
filter_scene=False,
filter_invalid=False,
ref_key="lmo_full",
),
lmo_bop_test=dict( lmo_bop_test=dict(
name="lmo_bop_test", name="lmo_bop_test",
dataset_root=osp.join(DATASETS_ROOT, "BOP_DATASETS/lmo/"), dataset_root=osp.join(DATASETS_ROOT, "BOP_DATASETS/lmo/test"),
models_root=osp.join(DATASETS_ROOT, "BOP_DATASETS/lmo/models"), models_root=osp.join(DATASETS_ROOT, "BOP_DATASETS/lmo/models"),
objs=LM_OCC_OBJECTS, objs=ref.lmo_full.objects,
ann_files=[osp.join(DATASETS_ROOT, "BOP_DATASETS/lmo/image_set/lmo_bop_test.txt")], ann_file=osp.join(DATASETS_ROOT, "BOP_DATASETS/lmo/test_targets_bop19.json"),
# NOTE: scene root
image_prefixes=[osp.join(DATASETS_ROOT, "BOP_DATASETS/lmo/test/{:06d}").format(2)],
xyz_prefixes=[None],
scale_to_meter=0.001, scale_to_meter=0.001,
with_masks=True, # (load masks but may not use it) with_masks=True, # (load masks but may not use it)
with_depth=True, # (load depth path here, but may not use it) with_depth=True, # (load depth path here, but may not use it)
...@@ -514,160 +328,31 @@ SPLITS_LM = dict( ...@@ -514,160 +328,31 @@ SPLITS_LM = dict(
cache_dir=osp.join(PROJ_ROOT, ".cache"), cache_dir=osp.join(PROJ_ROOT, ".cache"),
use_cache=True, use_cache=True,
num_to_load=-1, num_to_load=-1,
filter_scene=False,
filter_invalid=False, filter_invalid=False,
ref_key="lmo_full", ref_key="lmo_full",
), ),
) )
# single obj splits for lm real # single obj splits for lmo bop test
for obj in ref.lm_full.objects: for obj in ref.lmo_full.objects:
for split in ["train", "test", "all"]: for split in [
name = "lm_real_{}_{}".format(obj, split) "bop_test",
]:
name = "lmo_{}_{}".format(obj, split)
ann_files = [ ann_files = [
osp.join( osp.join(
DATASETS_ROOT, DATASETS_ROOT,
"BOP_DATASETS/lm/image_set/{}_{}.txt".format(obj, split), "BOP_DATASETS/lmo/image_set/{}_{}.txt".format(obj, split),
) )
] ]
if split in ["train", "all"]: # all is used to train lmo if name not in SPLITS_LMO:
filter_invalid = True SPLITS_LMO[name] = dict(
elif split in ["test"]:
filter_invalid = False
else:
raise ValueError("{}".format(split))
if name not in SPLITS_LM:
SPLITS_LM[name] = dict(
name=name,
dataset_root=osp.join(DATASETS_ROOT, "BOP_DATASETS/lm/"),
models_root=osp.join(DATASETS_ROOT, "BOP_DATASETS/lm/models"),
objs=[obj], # only this obj
ann_files=ann_files,
image_prefixes=[osp.join(DATASETS_ROOT, "BOP_DATASETS/lm/test/{:06d}").format(ref.lm_full.obj2id[obj])],
xyz_prefixes=[
osp.join(
DATASETS_ROOT,
"BOP_DATASETS/lm/test/xyz_crop/{:06d}".format(ref.lm_full.obj2id[obj]),
)
],
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,
num_to_load=-1,
filter_invalid=filter_invalid,
filter_scene=True,
ref_key="lm_full",
)
# single obj splits for lmo_NoBopTest_train
for obj in ref.lmo_full.objects:
for split in ["train"]:
name = "lmo_NoBopTest_{}_{}".format(obj, split)
if split in ["train"]:
filter_invalid = True
else:
raise ValueError("{}".format(split))
if name not in SPLITS_LM:
SPLITS_LM[name] = dict(
name=name,
dataset_root=osp.join(DATASETS_ROOT, "BOP_DATASETS/lmo/"),
models_root=osp.join(DATASETS_ROOT, "BOP_DATASETS/lmo/models"),
objs=[obj],
ann_files=[
osp.join(
DATASETS_ROOT,
"BOP_DATASETS/lmo/image_set/lmo_no_bop_test.txt",
)
],
# NOTE: scene root
image_prefixes=[osp.join(DATASETS_ROOT, "BOP_DATASETS/lmo/test/{:06d}").format(2)],
xyz_prefixes=[
osp.join(
DATASETS_ROOT,
"BOP_DATASETS/lmo/test/xyz_crop/{:06d}".format(2),
)
],
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,
num_to_load=-1,
filter_scene=False,
filter_invalid=filter_invalid,
ref_key="lmo_full",
)
# single obj splits for lmo_test
for obj in ref.lmo_full.objects:
for split in ["test"]:
name = "lmo_{}_{}".format(obj, split)
if split in ["train", "all"]: # all is used to train lmo
filter_invalid = True
elif split in ["test"]:
filter_invalid = False
else:
raise ValueError("{}".format(split))
if name not in SPLITS_LM:
SPLITS_LM[name] = dict(
name=name, name=name,
dataset_root=osp.join(DATASETS_ROOT, "BOP_DATASETS/lmo/"), dataset_root=osp.join(DATASETS_ROOT, "BOP_DATASETS/lmo/"),
models_root=osp.join(DATASETS_ROOT, "BOP_DATASETS/lmo/models"), models_root=osp.join(DATASETS_ROOT, "BOP_DATASETS/lmo/models"),
objs=[obj], objs=[obj], # only this obj
ann_files=[
osp.join(
DATASETS_ROOT,
"BOP_DATASETS/lmo/image_set/lmo_test.txt",
)
],
# NOTE: scene root
image_prefixes=[osp.join(DATASETS_ROOT, "BOP_DATASETS/lmo/test/{:06d}").format(2)],
xyz_prefixes=[None],
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,
num_to_load=-1,
filter_scene=False,
filter_invalid=False,
ref_key="lmo_full",
)
# single obj splits for lmo_bop_test
for obj in ref.lmo_full.objects:
for split in ["test"]:
name = "lmo_{}_bop_{}".format(obj, split)
if split in ["train", "all"]: # all is used to train lmo
filter_invalid = True
elif split in ["test"]:
filter_invalid = False
else:
raise ValueError("{}".format(split))
if name not in SPLITS_LM:
SPLITS_LM[name] = dict(
name=name,
dataset_root=osp.join(DATASETS_ROOT, "BOP_DATASETS/lmo/"),
models_root=osp.join(DATASETS_ROOT, "BOP_DATASETS/lmo/models"),
objs=[obj],
ann_files=[
osp.join(
DATASETS_ROOT,
"BOP_DATASETS/lmo/image_set/lmo_bop_test.txt",
)
],
# NOTE: scene root
image_prefixes=[osp.join(DATASETS_ROOT, "BOP_DATASETS/lmo/test/{:06d}").format(2)],
xyz_prefixes=[None],
scale_to_meter=0.001, scale_to_meter=0.001,
ann_file=osp.join(DATASETS_ROOT, "BOP_DATASETS/lmo/test_targets_bop19.json"),
with_masks=True, # (load masks but may not use it) with_masks=True, # (load masks but may not use it)
with_depth=True, # (load depth path here, but may not use it) with_depth=True, # (load depth path here, but may not use it)
height=480, height=480,
...@@ -675,60 +360,10 @@ for obj in ref.lmo_full.objects: ...@@ -675,60 +360,10 @@ for obj in ref.lmo_full.objects:
cache_dir=osp.join(PROJ_ROOT, ".cache"), cache_dir=osp.join(PROJ_ROOT, ".cache"),
use_cache=True, use_cache=True,
num_to_load=-1, num_to_load=-1,
filter_scene=False,
filter_invalid=False, filter_invalid=False,
ref_key="lmo_full", ref_key="lmo_full",
) )
# ================ add single image dataset for debug =======================================
debug_im_ids = {
"train": {obj: [] for obj in ref.lm_full.objects},
"test": {obj: [] for obj in ref.lm_full.objects},
}
for obj in ref.lm_full.objects:
for split in ["train", "test"]:
cur_ann_file = osp.join(DATASETS_ROOT, f"BOP_DATASETS/lm/image_set/{obj}_{split}.txt")
ann_files = [cur_ann_file]
im_ids = []
with open(cur_ann_file, "r") as f:
for line in f:
# scene_id(obj_id)/im_id
im_ids.append("{}/{}".format(ref.lm_full.obj2id[obj], int(line.strip("\r\n"))))
debug_im_ids[split][obj] = im_ids
for debug_im_id in debug_im_ids[split][obj]:
name = "lm_single_{}{}_{}".format(obj, debug_im_id.split("/")[1], split)
if name not in SPLITS_LM:
SPLITS_LM[name] = dict(
name=name,
dataset_root=osp.join(DATASETS_ROOT, "BOP_DATASETS/lm/"),
models_root=osp.join(DATASETS_ROOT, "BOP_DATASETS/lm/models"),
objs=[obj], # only this obj
ann_files=ann_files,
image_prefixes=[
osp.join(DATASETS_ROOT, "BOP_DATASETS/lm/test/{:06d}").format(ref.lm_full.obj2id[obj])
],
xyz_prefixes=[
osp.join(
DATASETS_ROOT,
"BOP_DATASETS/lm/test/xyz_crop/{:06d}".format(ref.lm_full.obj2id[obj]),
)
],
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,
num_to_load=-1,
filter_invalid=False,
filter_scene=True,
ref_key="lm_full",
debug_im_id=debug_im_id, # NOTE: debug im id
)
def register_with_name_cfg(name, data_cfg=None): def register_with_name_cfg(name, data_cfg=None):
"""Assume pre-defined datasets live in `./datasets`. """Assume pre-defined datasets live in `./datasets`.
...@@ -740,25 +375,25 @@ def register_with_name_cfg(name, data_cfg=None): ...@@ -740,25 +375,25 @@ def register_with_name_cfg(name, data_cfg=None):
data_cfg can be set in cfg.DATA_CFG.name data_cfg can be set in cfg.DATA_CFG.name
""" """
dprint("register dataset: {}".format(name)) dprint("register dataset: {}".format(name))
if name in SPLITS_LM: if name in SPLITS_LMO:
used_cfg = SPLITS_LM[name] used_cfg = SPLITS_LMO[name]
else: else:
assert data_cfg is not None, f"dataset name {name} is not registered" assert data_cfg is not None, f"dataset name {name} is not registered"
used_cfg = data_cfg used_cfg = data_cfg
DatasetCatalog.register(name, LM_Dataset(used_cfg)) DatasetCatalog.register(name, LMO_BOP_TEST_Dataset(used_cfg))
# something like eval_types # something like eval_types
MetadataCatalog.get(name).set( MetadataCatalog.get(name).set(
id="linemod", # NOTE: for pvnet to determine module id="lmo", # NOTE: for pvnet to determine module
ref_key=used_cfg["ref_key"], ref_key=used_cfg["ref_key"],
objs=used_cfg["objs"], objs=used_cfg["objs"],
eval_error_types=["ad", "rete", "proj"], eval_error_types=["ad", "rete", "proj"],
evaluator_type="bop", evaluator_type="bop",
**get_lm_metadata(obj_names=used_cfg["objs"], ref_key=used_cfg["ref_key"]), **get_lmo_metadata(obj_names=used_cfg["objs"], ref_key=used_cfg["ref_key"]),
) )
def get_available_datasets(): def get_available_datasets():
return list(SPLITS_LM.keys()) return list(SPLITS_LMO.keys())
#### tests ############################################### #### tests ###############################################
......
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment