BasicDataset returns a pair of image and labels (targets). If targets are not given, BasicDataset returns None as the label. This class supports strong augmentation for Fixmatch, and return both weakly and strongly augmented images.
| 10 | |
| 11 | |
| 12 | class BasicDataset(Dataset): |
| 13 | """ |
| 14 | BasicDataset returns a pair of image and labels (targets). |
| 15 | If targets are not given, BasicDataset returns None as the label. |
| 16 | This class supports strong augmentation for Fixmatch, |
| 17 | and return both weakly and strongly augmented images. |
| 18 | """ |
| 19 | |
| 20 | def __init__(self, |
| 21 | alg, |
| 22 | data, |
| 23 | targets=None, |
| 24 | num_classes=None, |
| 25 | transform=None, |
| 26 | is_ulb=False, |
| 27 | strong_transform=None, |
| 28 | onehot=False, |
| 29 | *args, **kwargs): |
| 30 | """ |
| 31 | Args |
| 32 | data: x_data |
| 33 | targets: y_data (if not exist, None) |
| 34 | num_classes: number of label classes |
| 35 | transform: basic transformation of data |
| 36 | use_strong_transform: If True, this dataset returns both weakly and strongly augmented images. |
| 37 | strong_transform: list of transformation functions for strong augmentation |
| 38 | onehot: If True, label is converted into onehot vector. |
| 39 | """ |
| 40 | super(BasicDataset, self).__init__() |
| 41 | self.alg = alg |
| 42 | self.data = data |
| 43 | self.targets = targets |
| 44 | |
| 45 | self.num_classes = num_classes |
| 46 | self.is_ulb = is_ulb |
| 47 | self.onehot = onehot |
| 48 | |
| 49 | self.transform = transform |
| 50 | if self.is_ulb: |
| 51 | if strong_transform is None: |
| 52 | self.strong_transform = copy.deepcopy(transform) |
| 53 | self.strong_transform.transforms.insert(0, RandAugment(3, 5)) |
| 54 | else: |
| 55 | self.strong_transform = strong_transform |
| 56 | |
| 57 | def __getitem__(self, idx): |
| 58 | """ |
| 59 | If strong augmentation is not used, |
| 60 | return weak_augment_image, target |
| 61 | else: |
| 62 | return weak_augment_image, strong_augment_image, target |
| 63 | """ |
| 64 | |
| 65 | # set idx-th target |
| 66 | if self.targets is None: |
| 67 | target = None |
| 68 | else: |
| 69 | target_ = self.targets[idx] |
no outgoing calls
no test coverage detected