| 118 | |
| 119 | |
| 120 | class RandomFlip: |
| 121 | def __init__(self, p): |
| 122 | # p = [p_x, p_y, p_z] probability of flipping each axis |
| 123 | assert len(p) == 3 |
| 124 | assert 0 < sum(p) <= 1, 'sum(p) must be in (0, 1] range, is: {}'.format(sum(p)) |
| 125 | self.p = p |
| 126 | self.p_cum_sum = np.cumsum(p) |
| 127 | |
| 128 | def __call__(self, coords): |
| 129 | r = random.random() |
| 130 | if r <= self.p_cum_sum[0]: |
| 131 | # Flip the first axis |
| 132 | coords[..., 0] = -coords[..., 0] |
| 133 | elif r <= self.p_cum_sum[1]: |
| 134 | # Flip the second axis |
| 135 | coords[..., 1] = -coords[..., 1] |
| 136 | elif r <= self.p_cum_sum[2]: |
| 137 | # Flip the third axis |
| 138 | coords[..., 2] = -coords[..., 2] |
| 139 | |
| 140 | return coords |
| 141 | |
| 142 | |
| 143 | class RandomRotation: |
nothing calls this directly
no outgoing calls
no test coverage detected