| 12 | |
| 13 | |
| 14 | class EventAugment(object): |
| 15 | def __init__(self, resolution): |
| 16 | self.resolution = resolution |
| 17 | self.augment_list = [ |
| 18 | (self.identity, 0, 0), |
| 19 | (self.drop_by_time, 0.1, 0.9), |
| 20 | (self.drop_by_area, 0.1, 0.5), |
| 21 | (self.random_drop, 0.1, 0.5), |
| 22 | # (self.drop_by_area_with_cam, 0.1, 0.6), |
| 23 | # (self.random_drop_with_cam, 0.5, 1), |
| 24 | (self.overall_noise, 0.1, 0.9), |
| 25 | (self.region_noise, 0.1, 0.5), |
| 26 | # (self.overall_noise_with_cam, 0.1, 1), |
| 27 | # (self.region_noise_with_cam, 0.1, 0.9), |
| 28 | (self.time_incline_x, 0.05, 0.5), |
| 29 | (self.time_incline_y, 0.05, 0.5), |
| 30 | # (self.random_shift_time, 0.1, 0.8), |
| 31 | |
| 32 | (self.random_shift_xy, 1, 10), |
| 33 | (self.flip_along_x, 0, 0), |
| 34 | (self.flip_along_y, 0, 0), |
| 35 | (self.flip_along_time, 0, 0), |
| 36 | (self.rotate, 0, math.pi / 2), |
| 37 | (self.linear_x, 0, 0.6), |
| 38 | (self.linear_y, 0, 0.6), |
| 39 | (self.shear_x, 0, 1), |
| 40 | (self.shear_y, 0, 1), |
| 41 | (self.scale, 0.2, 2)] |
| 42 | self.ops_name = [] |
| 43 | self.ops_list = [] |
| 44 | self.mags_list = [] |
| 45 | self.l_ops = len(self.augment_list) |
| 46 | self.l_uniq = 0 |
| 47 | for idx, op in enumerate(self.augment_list): |
| 48 | self.ops_name.append(op.__str__().split(' ')[2].split('.')[1]) |
| 49 | |
| 50 | def __call__(self, events): |
| 51 | op_idx = random.randint(0, len(self.augment_list)) - 1 |
| 52 | op_max = self.augment_list[op_idx][2] |
| 53 | op_min = self.augment_list[op_idx][1] |
| 54 | op = self.augment_list[op_idx][0] |
| 55 | aug_events = op(events, random.random() * (op_max - op_min) + op_min) |
| 56 | return aug_events |
| 57 | |
| 58 | def identity(self, events, v): |
| 59 | events = copy.deepcopy(events) |
| 60 | return events |
| 61 | |
| 62 | |
| 63 | |
| 64 | def overall_noise(self, events, ratio): |
| 65 | events = copy.deepcopy(events).to(events.device) |
| 66 | t_max = torch.amax(events[:, 2]).item() |
| 67 | t_min = torch.amin(events[:, 2]).item() |
| 68 | len_noise = int(len(events) * ratio) |
| 69 | x_noise = torch.randint(high=self.resolution[1], size=(len_noise, 1)) |
| 70 | y_noise = torch.randint(high=self.resolution[0], size=(len_noise, 1)) |
| 71 | t_noise = torch.rand(size=(len_noise, 1)) * (t_max - t_min) + t_min |