Spaces:
Sleeping
Sleeping
| import albumentations as albu | |
| import numpy as np | |
| import cv2 | |
| import os | |
| os.environ['CUDA_VISIBLE_DEVICES'] = '0' | |
| class Dataset: | |
| def __init__( | |
| self, | |
| image_path, | |
| augmentation=None, | |
| preprocessing=None, | |
| ): | |
| self.pil_image = image_path | |
| self.augmentation = augmentation | |
| self.preprocessing = preprocessing | |
| def get(self): | |
| # pil image > numpy array | |
| image = np.array(self.pil_image) | |
| image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) | |
| # apply augmentations | |
| if self.augmentation: | |
| sample = self.augmentation(image=image) | |
| image = sample['image'] | |
| # apply preprocessing | |
| if self.preprocessing: | |
| sample = self.preprocessing(image=image) | |
| image = sample['image'] | |
| return image | |
| def get_validation_augmentation(): | |
| """Add paddings to make image shape divisible by 32""" | |
| test_transform = [ | |
| albu.PadIfNeeded(384, 480) | |
| ] | |
| return albu.Compose(test_transform) | |
| def to_tensor(x, **kwargs): | |
| return x.transpose(2, 0, 1).astype('float32') | |
| def get_preprocessing(preprocessing_fn): | |
| _transform = [ | |
| albu.Lambda(image=preprocessing_fn), | |
| albu.Lambda(image=to_tensor), | |
| ] | |
| return albu.Compose(_transform) | |