52 lines
1.6 KiB
Python
52 lines
1.6 KiB
Python
import albumentations as a
|
|
import numpy as np
|
|
from albumentations.pytorch import ToTensorV2
|
|
from hyperparams import CROP_SIZE
|
|
|
|
|
|
preprocess_training = a.Compose(
|
|
[
|
|
a.augmentations.PadIfNeeded(min_width=CROP_SIZE, min_height=CROP_SIZE),
|
|
a.RandomCrop(width=CROP_SIZE, height=CROP_SIZE),
|
|
a.GaussNoise(),
|
|
a.Flip(p=0.5),
|
|
a.RandomRotate90(p=0.5),
|
|
a.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
|
|
ToTensorV2(),
|
|
]
|
|
)
|
|
preprocess_validation = a.Compose(
|
|
[
|
|
a.augmentations.PadIfNeeded(min_width=CROP_SIZE, min_height=CROP_SIZE),
|
|
a.CenterCrop(width=CROP_SIZE, height=CROP_SIZE),
|
|
a.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
|
|
ToTensorV2(),
|
|
]
|
|
)
|
|
|
|
|
|
def transform_training(example):
|
|
transformed = []
|
|
for pil_image in example["image"]:
|
|
array = np.array(pil_image.convert("RGB"))
|
|
# check if image is in (height, width, channel) shape
|
|
# if not, do a transpose
|
|
if array.shape[-1] != 3:
|
|
array = np.transpose(array, (1, 2, 0))
|
|
img = preprocess_training(image=array)["image"]
|
|
transformed.append(img)
|
|
example["pixel_values"] = transformed
|
|
return example
|
|
|
|
|
|
def transform_validation(example):
|
|
transformed = []
|
|
for pil_image in example["image"]:
|
|
array = np.array(pil_image.convert("RGB"))
|
|
if array.shape[-1] != 3:
|
|
array = np.transpose(array, (1, 2, 0))
|
|
img = preprocess_validation(image=array)["image"]
|
|
transformed.append(img)
|
|
example["pixel_values"] = transformed
|
|
return example
|