(dataset, include_classes=range(160))
| 145 | |
| 146 | |
| 147 | def subsample_classes(dataset, include_classes=range(160)): |
| 148 | |
| 149 | include_classes_cub = np.array(include_classes) + 1 # CUB classes are indexed 1 --> 200 instead of 0 --> 199 |
| 150 | # cls_idxs = [x for x, (_, r) in enumerate(dataset.data.iterrows()) if int(r['target']) in include_classes_cub] |
| 151 | cls_idxs = [x for x, r in enumerate(dataset.data) if int(r[2]) in include_classes_cub] |
| 152 | |
| 153 | # TODO: For now have no target transform |
| 154 | target_xform_dict = {} |
| 155 | for i, k in enumerate(include_classes): |
| 156 | target_xform_dict[k] = i |
| 157 | |
| 158 | dataset = subsample_dataset(dataset, cls_idxs) |
| 159 | |
| 160 | dataset.target_transform = lambda x: target_xform_dict[x] |
| 161 | |
| 162 | return dataset |
| 163 | |
| 164 | |
| 165 | def get_train_val_indices(train_dataset, val_split=0.2): |
no test coverage detected