def __init__(self, directory=None, split=None, exp_dict=None):
    self.class_ids = CLASS_IDS
    n_classes = exp_dict.get('n_classes')
    if n_classes:
        self.class_ids = CLASS_IDS[:n_classes]
    classes = exp_dict.get('classes')
    if classes:
        classes_map = {key: value for (value, key) in class_ids_map.items()}
        self.class_ids = sorted([classes_map[k] for k in classes])
    self.split = split
    self.elevation = 30.0
    self.distance = 2.732
    self.exp_dict = exp_dict
    self.class_ids_map = class_ids_map
    self.images = []
    self.voxels = []
    self.labels = []
    self.class_ids_pair = list(zip(self.class_ids, [self.class_ids_map[i] for i in self.class_ids]))
    self.num_data = {}
    self.pos = {}
    count = 0
    loop = tqdm.tqdm(self.class_ids)
    loop.set_description(f'Loading {split} Dataset')
    n_train_objects = exp_dict.get('n_train_objects')
    n_ratio_val = exp_dict.get('n_val_ratio')
    if ((n_train_objects is None) and (split == 'unlabeled')):
        return
    if (split in ['train', 'unlabeled']):
        set_name = 'train'
    elif (split in ['val', 'test']):
        set_name = 'val'
        if (n_ratio_val is None):
            set_name = split
    for (ci, class_id) in enumerate(loop):
        i = list(np.load(os.path.join(directory, ('%s_%s_images.npz' % (class_id, set_name)))).items())[0][1]
        v = list(np.load(os.path.join(directory, ('%s_%s_voxels.npz' % (class_id, set_name)))).items())[0][1]
        if ((split == 'train') and (n_train_objects is not None)):
            n = n_train_objects
            i = i[:n]
            v = v[:n]
        if ((split == 'unlabeled') and (n_train_objects is not None)):
            n = n_train_objects
            i = i[n:]
            v = v[n:]
        elif ((split == 'val') and (n_ratio_val is not None)):
            n = int((i.shape[0] * n_ratio_val))
            i = i[:n]
            v = v[:n]
        elif ((split == 'test') and (n_ratio_val is not None)):
            n = int((i.shape[0] * n_ratio_val))
            i = i[n:]
            v = v[n:]
        self.images += [i]
        self.voxels += [v]
        self.labels += [(torch.ones(i.shape[0]) * ci)]
    self.images = np.concatenate(self.images, axis=0)
    self.images = torch.from_numpy((self.images.astype('float32') / 255.0))
    self.voxels = np.concatenate(self.voxels, axis=0)
    self.voxels = torch.from_numpy(self.voxels.astype('float32'))
    self.labels = torch.cat(self.labels, dim=0)
    distances = (torch.ones(24).float() * self.distance)
    elevations = (torch.ones(24).float() * self.elevation)
    self.possible_viewpoints = srf.get_points_from_angles(distances, elevations, ((- torch.arange(24)) * 15))
    print(f'{split} samples: {len(self)}')
def __init__(self, directory=None, split=None, exp_dict=None):
    self.class_ids = CLASS_IDS
    n_classes = exp_dict.get('n_classes')
    if n_classes:
        self.class_ids = CLASS_IDS[:n_classes]
    classes = exp_dict.get('classes')
    if classes:
        classes_map = {key: value for (value, key) in class_ids_map.items()}
        self.class_ids = sorted([classes_map[k] for k in classes])
    self.split = split
    self.elevation = 30.0
    self.distance = 2.732
    self.exp_dict = exp_dict
    self.class_ids_map = class_ids_map
    self.images = []
    self.voxels = []
    self.labels = []
    self.class_ids_pair = list(zip(self.class_ids, [self.class_ids_map[i] for i in self.class_ids]))
    self.num_data = {}
    self.pos = {}
    count = 0
    loop = tqdm.tqdm(self.class_ids)
    loop.set_description(f'Loading {split} Dataset')
    n_train_objects = exp_dict.get('n_train_objects')
    n_ratio_val = exp_dict.get('n_val_ratio')
    if ((n_train_objects is None) and (split == 'unlabeled')):
        return
    if (split in ['train', 'unlabeled']):
        set_name = 'train'
    elif (split in ['val', 'test']):
        set_name = 'val'
        if (n_ratio_val is None):
            set_name = split
    for (ci, class_id) in enumerate(loop):
        i = list(np.load(os.path.join(directory, ('%s_%s_images.npz' % (class_id, set_name)))).items())[0][1]
        v = list(np.load(os.path.join(directory, ('%s_%s_voxels.npz' % (class_id, set_name)))).items())[0][1]
        if ((split == 'train') and (n_train_objects is not None)):
            n = n_train_objects
            i = i[:n]
            v = v[:n]
        if ((split == 'unlabeled') and (n_train_objects is not None)):
            n = n_train_objects
            i = i[n:]
            v = v[n:]
        elif ((split == 'val') and (n_ratio_val is not None)):
            n = int((i.shape[0] * n_ratio_val))
            i = i[:n]
            v = v[:n]
        elif ((split == 'test') and (n_ratio_val is not None)):
            n = int((i.shape[0] * n_ratio_val))
            i = i[n:]
            v = v[n:]
        self.images += [i]
        self.voxels += [v]
        self.labels += [(torch.ones(i.shape[0]) * ci)]
    self.images = np.concatenate(self.images, axis=0)
    self.images = torch.from_numpy((self.images.astype('float32') / 255.0))
    self.voxels = np.concatenate(self.voxels, axis=0)
    self.voxels = torch.from_numpy(self.voxels.astype('float32'))
    self.labels = torch.cat(self.labels, dim=0)
    distances = (torch.ones(24).float() * self.distance)
    elevations = (torch.ones(24).float() * self.elevation)
    self.possible_viewpoints = srf.get_points_from_angles(distances, elevations, ((- torch.arange(24)) * 15))
    print(f'{split} samples: {len(self)}')