| 51 | |
| 52 | |
| 53 | class SequentialShapeRandomPointcloudPatchSampler(data.sampler.Sampler): |
| 54 | |
| 55 | def __init__(self, data_source, patches_per_shape, seed=None, sequential_shapes=False, identical_epochs=False): |
| 56 | self.data_source = data_source |
| 57 | self.patches_per_shape = patches_per_shape |
| 58 | self.sequential_shapes = sequential_shapes |
| 59 | self.seed = seed |
| 60 | self.identical_epochs = identical_epochs |
| 61 | self.total_patch_count = None |
| 62 | self.shape_patch_inds = None |
| 63 | |
| 64 | if self.seed is None: |
| 65 | self.seed = np.random.random_integers(0, 2**32-1, 1)[0] |
| 66 | self.rng = np.random.RandomState(self.seed) |
| 67 | |
| 68 | self.total_patch_count = 0 |
| 69 | for shape_ind, _ in enumerate(self.data_source.shape_names): |
| 70 | self.total_patch_count = self.total_patch_count + min(self.patches_per_shape, self.data_source.shape_patch_count[shape_ind]) |
| 71 | |
| 72 | def __iter__(self): |
| 73 | |
| 74 | # optionally always pick the same permutation (mainly for debugging) |
| 75 | if self.identical_epochs: |
| 76 | self.rng.seed(self.seed) |
| 77 | |
| 78 | # global point index offset for each shape |
| 79 | shape_patch_offset = list(np.cumsum(self.data_source.shape_patch_count)) |
| 80 | shape_patch_offset.insert(0, 0) |
| 81 | shape_patch_offset.pop() |
| 82 | |
| 83 | shape_inds = range(len(self.data_source.shape_names)) |
| 84 | |
| 85 | if not self.sequential_shapes: |
| 86 | shape_inds = self.rng.permutation(shape_inds) |
| 87 | |
| 88 | # return a permutation of the points in the dataset where all points in the same shape are adjacent (for performance reasons): |
| 89 | # first permute shapes, then concatenate a list of permuted points in each shape |
| 90 | self.shape_patch_inds = [[]]*len(self.data_source.shape_names) |
| 91 | point_permutation = [] |
| 92 | for shape_ind in shape_inds: |
| 93 | start = shape_patch_offset[shape_ind] |
| 94 | end = shape_patch_offset[shape_ind]+self.data_source.shape_patch_count[shape_ind] |
| 95 | |
| 96 | global_patch_inds = self.rng.choice(range(start, end), size=min(self.patches_per_shape, end-start), replace=False) |
| 97 | point_permutation.extend(global_patch_inds) |
| 98 | |
| 99 | # save indices of shape point subset |
| 100 | self.shape_patch_inds[shape_ind] = global_patch_inds - start |
| 101 | |
| 102 | return iter(point_permutation) |
| 103 | |
| 104 | def __len__(self): |
| 105 | return self.total_patch_count |
| 106 | |
| 107 | class RandomPointcloudPatchSampler(data.sampler.Sampler): |
| 108 |
no outgoing calls