Spaces:
Sleeping
Sleeping
File size: 1,353 Bytes
c992a7d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 |
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)
|