Class MeshConcatDataset

pointersect/data/mesh_dataset_v2.py:376–419

r"""Dataset as a concatenation of multiple MeshDatasets. Arguments: datasets (sequence): List of datasets to be concatenated

class MeshConcatDataset(torch.utils.data.Dataset):
    r"""Dataset as a concatenation of multiple MeshDatasets.

    Arguments:
        datasets (sequence): List of datasets to be concatenated
    """

    def __init__(
            self,
            datasets: T.Iterable[MeshDataset],
            max_retry=30,
            wait_sec=5,
    ):

        self.datasets = list(datasets)
        assert len(self.datasets) > 0
        self.concat_dataset = torch.utils.data.ConcatDataset(self.datasets)
        self.max_retry = max_retry
        self.wait_sec = wait_sec

    def __len__(self):
        return len(self.concat_dataset)

    def __getitem__(self, idx):
        d = None
        for retry in range(self.max_retry):
            try:
                d = self.concat_dataset[idx]
                if d is not None:
                    return d
            except:
                traceback.print_exc()
                time.sleep(self.wait_sec)
        return d

    def get_all_num_pixels(self):
        all_seq_lens = []
        for dset in self.datasets:
            seq_lens = dset.get_all_num_pixels()
            if seq_lens is None:
                return None
            else:
                all_seq_lens += seq_lens
        return all_seq_lens