Spaces:
Runtime error
Runtime error
File size: 3,005 Bytes
056188b d7c1491 056188b 4cf5013 056188b 4cf5013 056188b 4cf5013 056188b d7c1491 4cf5013 056188b |
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 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 |
import tensorflow as tf
from typing import List
from ..commons import read_image
from ..augmentation import UnpairedAugmentationFactory
class UnpairedLowLightDataset:
def __init__(
self,
image_size: int = 256,
apply_resize: bool = False,
apply_random_horizontal_flip: bool = True,
apply_random_vertical_flip: bool = True,
apply_random_rotation: bool = True,
) -> None:
self.augmentation_factory = UnpairedAugmentationFactory(image_size=image_size)
self.image_size = image_size
self.apply_resize = apply_resize
self.apply_random_horizontal_flip = apply_random_horizontal_flip
self.apply_random_vertical_flip = apply_random_vertical_flip
self.apply_random_rotation = apply_random_rotation
def _resize(self, image):
return tf.image.resize(image, (self.image_size, self.image_size))
def _get_dataset(self, images: List[str], batch_size: int, is_train: bool):
dataset = tf.data.Dataset.from_tensor_slices((images))
dataset = dataset.map(read_image, num_parallel_calls=tf.data.AUTOTUNE)
dataset = (
dataset.map(
self.augmentation_factory.random_crop,
num_parallel_calls=tf.data.AUTOTUNE,
)
if not self.apply_resize
else dataset.map(self._resize, num_parallel_calls=tf.data.AUTOTUNE)
)
if is_train:
dataset = (
dataset.map(
self.augmentation_factory.random_horizontal_flip,
num_parallel_calls=tf.data.AUTOTUNE,
)
if self.apply_random_horizontal_flip
else dataset
)
dataset = (
dataset.map(
self.augmentation_factory.random_vertical_flip,
num_parallel_calls=tf.data.AUTOTUNE,
)
if self.apply_random_vertical_flip
else dataset
)
dataset = (
dataset.map(
self.augmentation_factory.random_rotate,
num_parallel_calls=tf.data.AUTOTUNE,
)
if self.apply_random_rotation
else dataset
)
dataset = dataset.batch(batch_size, drop_remainder=True)
return dataset
def get_datasets(
self,
images: List[str],
val_split: float = 0.2,
batch_size: int = 16,
):
split_index = int(len(images) * (1 - val_split))
train_images = images[:split_index]
val_images = images[split_index:]
print(f"Number of train data points: {len(train_images)}")
print(f"Number of validation data points: {len(val_images)}")
train_dataset = self._get_dataset(train_images, batch_size, is_train=True)
val_dataset = self._get_dataset(val_images, batch_size, is_train=False)
return train_dataset, val_dataset
|