Beispiel #1
0
    def get_data_transforms(self):
        cfg = self.cfg
        bbox_params = self.get_bbox_params()
        transform = Compose([Resize(cfg.data.img_sz, cfg.data.img_sz)],
                            bbox_params=bbox_params)

        augmentors_dict = {
            'Blur': Blur(),
            'RandomRotate90': RandomRotate90(),
            'HorizontalFlip': HorizontalFlip(),
            'VerticalFlip': VerticalFlip(),
            'GaussianBlur': GaussianBlur(),
            'GaussNoise': GaussNoise(),
            'RGBShift': RGBShift(),
            'ToGray': ToGray()
        }
        aug_transforms = []
        for augmentor in cfg.data.augmentors:
            try:
                aug_transforms.append(augmentors_dict[augmentor])
            except KeyError as e:
                log.warning(
                    '{0} is an unknown augmentor. Continuing without {0}. \
                    Known augmentors are: {1}'.format(
                        e, list(augmentors_dict.keys())))
        aug_transforms.append(Resize(cfg.data.img_sz, cfg.data.img_sz))
        aug_transform = Compose(aug_transforms, bbox_params=bbox_params)

        return transform, aug_transform
Beispiel #2
0
    def get_data_transforms(self) -> Tuple[BasicTransform, BasicTransform]:
        """Get albumentations transform objects for data augmentation.

        Returns:
           1st tuple arg: a transform that doesn't do any data augmentation
           2nd tuple arg: a transform with data augmentation
        """
        cfg = self.cfg
        bbox_params = self.get_bbox_params()
        transform = Compose([Resize(cfg.data.img_sz, cfg.data.img_sz)],
                            bbox_params=bbox_params)

        augmentors_dict = {
            'Blur': Blur(),
            'RandomRotate90': RandomRotate90(),
            'HorizontalFlip': HorizontalFlip(),
            'VerticalFlip': VerticalFlip(),
            'GaussianBlur': GaussianBlur(),
            'GaussNoise': GaussNoise(),
            'RGBShift': RGBShift(),
            'ToGray': ToGray()
        }
        aug_transforms = []
        for augmentor in cfg.data.augmentors:
            try:
                aug_transforms.append(augmentors_dict[augmentor])
            except KeyError as e:
                log.warning(
                    '{0} is an unknown augmentor. Continuing without {0}. \
                    Known augmentors are: {1}'.format(
                        e, list(augmentors_dict.keys())))
        aug_transforms.append(Resize(cfg.data.img_sz, cfg.data.img_sz))
        aug_transform = Compose(aug_transforms, bbox_params=bbox_params)

        return transform, aug_transform
Beispiel #3
0
def build_databunch(data_dir, img_sz, batch_sz, class_names, augmentors):
    num_workers = 4

    train_dir = join(data_dir, 'train')
    valid_dir = join(data_dir, 'valid')

    augmentors_dict = {
        'Blur': Blur(),
        'RandomRotate90': RandomRotate90(),
        'HorizontalFlip': HorizontalFlip(),
        'VerticalFlip': VerticalFlip(),
        'GaussianBlur': GaussianBlur(),
        'GaussNoise': GaussNoise(),
        'RGBShift': RGBShift(),
        'ToGray': ToGray()
    }

    aug_transforms = []
    for augmentor in augmentors:
        try:
            aug_transforms.append(augmentors_dict[augmentor])
        except KeyError as e:
            log.warning('{0} is an unknown augmentor. Continuing without {0}. \
                Known augmentors are: {1}'.format(e,
                                                  list(
                                                      augmentors_dict.keys())))
    aug_transforms = Compose(aug_transforms)

    train_ds = AlbumentationDataset(ImageFolder(train_dir,
                                                classes=class_names),
                                    transform=aug_transforms)
    valid_ds = AlbumentationDataset(ImageFolder(valid_dir,
                                                classes=class_names))

    train_dl = DataLoader(train_ds,
                          shuffle=True,
                          batch_size=batch_sz,
                          num_workers=num_workers,
                          drop_last=True,
                          pin_memory=True)
    valid_dl = DataLoader(valid_ds,
                          batch_size=batch_sz,
                          num_workers=num_workers,
                          pin_memory=True)

    return DataBunch(train_ds, train_dl, valid_ds, valid_dl, class_names)
Beispiel #4
0
def build_databunch(data_dir, img_sz, batch_sz, class_names, rare_classes,
                    desired_prob, augmentors):
    num_workers = 4

    train_dir = join(data_dir, 'train')
    valid_dir = join(data_dir, 'valid')

    augmentors_dict = {
        'Blur': Blur(),
        'RandomRotate90': RandomRotate90(),
        'HorizontalFlip': HorizontalFlip(),
        'VerticalFlip': VerticalFlip(),
        'GaussianBlur': GaussianBlur(),
        'GaussNoise': GaussNoise(),
        'RGBShift': RGBShift(),
        'ToGray': ToGray()
    }

    aug_transforms = []
    for augmentor in augmentors:
        try:
            aug_transforms.append(augmentors_dict[augmentor])
        except KeyError as e:
            log.warning('{0} is an unknown augmentor. Continuing without {0}. \
                Known augmentors are: {1}'.format(e,
                                                  list(
                                                      augmentors_dict.keys())))
    aug_transforms = Compose(aug_transforms)

    train_ds = AlbumentationDataset(ImageFolder(train_dir,
                                                classes=class_names),
                                    transform=aug_transforms)
    valid_ds = AlbumentationDataset(ImageFolder(valid_dir,
                                                classes=class_names))

    if rare_classes != []:
        targets = [target for _, target in train_ds.orig_dataset.imgs]
        train_sample_weights = calculate_oversampling_weights(
            targets, rare_classes, desired_prob)
        num_train_samples = len(train_ds)
        train_sampler = WeightedRandomSampler(weights=train_sample_weights,
                                              num_samples=num_train_samples,
                                              replacement=True)
        shuffle = False
    else:
        train_sampler = None
        shuffle = True

    train_dl = DataLoader(train_ds,
                          shuffle=shuffle,
                          batch_size=batch_sz,
                          num_workers=num_workers,
                          drop_last=True,
                          pin_memory=True,
                          sampler=train_sampler)
    valid_dl = DataLoader(valid_ds,
                          batch_size=batch_sz,
                          num_workers=num_workers,
                          pin_memory=True)

    return DataBunch(train_ds, train_dl, valid_ds, valid_dl, class_names)