File size: 656 Bytes
5ce8761 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 | from .rlbench import (
Peract2Dataset,
Peract2SingleCamDataset,
PeractDataset,
PeractTwoCamDataset,
HiveformerDataset
)
def fetch_dataset_class(dataset_name):
"""Fetch the dataset class based on the dataset name."""
dataset_classes = {
"Peract2_3dfront_3dwrist": Peract2Dataset,
"Peract2_3dfront": Peract2SingleCamDataset,
"Peract": PeractDataset,
"PeractTwoCam": PeractTwoCamDataset,
"HiveformerRLBench": HiveformerDataset
}
if dataset_name not in dataset_classes:
raise ValueError(f"Unknown dataset: {dataset_name}")
return dataset_classes[dataset_name]
|