Функция image_dataset_from_directory

Функция image_dataset_from_directory создает датасет tf.data.Dataset из изображений, которые лежат в директории, разбитой на подпапки по классам. Первым параметром функция принимает путь к директории с изображениями. Вторым параметром можно передать размер изображений image_size, третьим - размер батча batch_size. Также можно указать labels, label_mode, color_mode, shuffle, seed, validation_split и subset.

Синтаксис

tf.keras.utils.image_dataset_from_directory(
	directory,
	labels='inferred',
	label_mode='int',
	class_names=None,
	color_mode='rgb',
	batch_size=32,
	image_size=(256, 256),
	shuffle=True,
	seed=None,
	validation_split=None,
	subset=None,
	interpolation='bilinear',
	follow_links=False,
)

Пример

Давайте создадим директорию с двумя классами изображений и загрузим из нее датасет:

import tensorflow as tf
import numpy as np
import os

tf.random.set_seed(0)

# create dummy images in two class folders
os.makedirs('data/cats', exist_ok=True)
os.makedirs('data/dogs', exist_ok=True)

for i in range(4):
	img = np.random.randint(0, 256, (64, 64, 3), dtype=np.uint8)
	tf.keras.utils.save_img(f'data/cats/{i}.png', img)
	img = np.random.randint(0, 256, (64, 64, 3), dtype=np.uint8)
	tf.keras.utils.save_img(f'data/dogs/{i}.png', img)

ds = tf.keras.utils.image_dataset_from_directory(
	'data',
	image_size=(32, 32),
	batch_size=2,
	seed=0,
)

print(ds)

Результат выполнения кода:

"Found 8 files belonging to 2 classes."

Пример

Давайте извлечем один батч изображений и меток из созданного датасета:

import tensorflow as tf

ds = tf.keras.utils.image_dataset_from_directory(
	'data',
	image_size=(32, 32),
	batch_size=2,
	seed=0,
)

for images, labels in ds.take(1):
	print(images.shape)
	print(labels)

Результат выполнения кода:

(2, 32, 32, 3)
tf.Tensor([0 1], shape=(2,), dtype=int32)

Пример

Давайте загрузим датасет в оттенках серого и с однотипными метками в формате one-hot:

import tensorflow as tf

ds = tf.keras.utils.image_dataset_from_directory(
	'data',
	image_size=(32, 32),
	batch_size=2,
	color_mode='grayscale',
	label_mode='categorical',
	seed=0,
)

for images, labels in ds.take(1):
	print(images.shape)
	print(labels)

Результат выполнения кода:

(2, 32, 32, 1)
tf.Tensor(
[[1. 0.]
 [0. 1.]], shape=(2, 2), dtype=float32)

Пример

Давайте разделим датасет на обучающую и валидационную части с помощью параметра validation_split:

import tensorflow as tf

train_ds = tf.keras.utils.image_dataset_from_directory(
	'data',
	image_size=(32, 32),
	batch_size=2,
	validation_split=0.25,
	subset='training',
	seed=0,
)

val_ds = tf.keras.utils.image_dataset_from_directory(
	'data',
	image_size=(32, 32),
	batch_size=2,
	validation_split=0.25,
	subset='validation',
	seed=0,
)

print(train_ds)
print(val_ds)

Результат выполнения кода:

"Found 8 files belonging to 2 classes. Using 6 files for training."
"Found 8 files belonging to 2 classes. Using 2 files for validation."

Смотрите также

функцию load_img,
которая загружает одно изображение из файла
функцию img_to_array,
которая преобразует изображение в массив NumPy
функцию array_to_img,
которая преобразует массив NumPy в изображение
функцию save_img,
которая сохраняет изображение в файл