| 39 | |
| 40 | |
| 41 | class NCaltech101(Dataset): |
| 42 | def __init__(self, data_path='data/n-caltech/frames_number_10_split_by_number', |
| 43 | data_type='train', transform=False): |
| 44 | |
| 45 | self.filepath = os.path.join(data_path) |
| 46 | self.clslist = os.listdir(self.filepath) |
| 47 | self.clslist.sort() |
| 48 | |
| 49 | self.dvs_filelist = [] |
| 50 | self.targets = [] |
| 51 | self.resize = transforms.Resize(size=(48, 48), interpolation=torchvision.transforms.InterpolationMode.NEAREST) |
| 52 | |
| 53 | for i, cls in enumerate(self.clslist): |
| 54 | # print (i, cls) |
| 55 | file_list = os.listdir(os.path.join(self.filepath, cls)) |
| 56 | num_file = len(file_list) |
| 57 | |
| 58 | cut_idx = int(num_file * 0.9) |
| 59 | train_file_list = file_list[:cut_idx] |
| 60 | test_split_list = file_list[cut_idx:] |
| 61 | for file in file_list: |
| 62 | if data_type == 'train': |
| 63 | if file in train_file_list: |
| 64 | self.dvs_filelist.append(os.path.join(self.filepath, cls, file)) |
| 65 | self.targets.append(i) |
| 66 | else: |
| 67 | if file in test_split_list: |
| 68 | self.dvs_filelist.append(os.path.join(self.filepath, cls, file)) |
| 69 | self.targets.append(i) |
| 70 | |
| 71 | self.data_num = len(self.dvs_filelist) |
| 72 | self.data_type = data_type |
| 73 | if data_type != 'train': |
| 74 | counts = np.unique(np.array(self.targets), return_counts=True)[1] |
| 75 | class_weights = counts.sum() / (counts * len(counts)) |
| 76 | self.class_weights = torch.Tensor(class_weights) |
| 77 | self.classes = range(101) |
| 78 | self.transform = transform |
| 79 | self.rotate = transforms.RandomRotation(degrees=15) |
| 80 | self.shearx = transforms.RandomAffine(degrees=0, shear=(-15, 15)) |
| 81 | |
| 82 | def __getitem__(self, index): |
| 83 | file_pth = self.dvs_filelist[index] |
| 84 | label = self.targets[index] |
| 85 | data = torch.from_numpy(np.load(file_pth)['frames']).float() |
| 86 | data = self.resize(data) |
| 87 | |
| 88 | if self.transform: |
| 89 | |
| 90 | choices = ['roll', 'rotate', 'shear'] |
| 91 | aug = np.random.choice(choices) |
| 92 | if aug == 'roll': |
| 93 | off1 = random.randint(-3, 3) |
| 94 | off2 = random.randint(-3, 3) |
| 95 | data = torch.roll(data, shifts=(off1, off2), dims=(2, 3)) |
| 96 | if aug == 'rotate': |
| 97 | data = self.rotate(data) |
| 98 | if aug == 'shear': |