Папка картинок в TensorFlow

Если снимки уже лежат на диске, их не обязательно собирать вручную в тензор. Функция image_dataset_from_directory обходит папку картинок: имя каждого подкаталога задаёт класс, файлы внутри - сами изображения.

Для примера создадим каталог mini_album с двумя подпапками классов и парой маленьких файлов, затем построим набор данных с фиксированным размером картинки:

import os
import tensorflow as tf

base = "mini_album"
for cls in ("cats", "dogs"):
	folder = os.path.join(base, cls)
	os.makedirs(folder, exist_ok=True)
	for i in range(2):
		img = tf.zeros((8, 8, 3), dtype=tf.float32)
		if cls == "cats":
			img = img + tf.constant([[[1.0, 0.0, 0.0]]])
		else:
			img = img + tf.constant([[[0.0, 0.0, 1.0]]])
		img = tf.cast(img * 255, tf.uint8)
		path = os.path.join(folder, f"{i}.png")
		tf.io.write_file(path, tf.io.encode_png(img))

ds = tf.keras.utils.image_dataset_from_directory(
	base,
	image_size=(8, 8),
	batch_size=2,
	shuffle=False,
)
for images, labels in ds.take(1):
	print(images.shape)
	print(labels)
	# выведет (2, 8, 8, 3)

Первый аргумент - путь к корню, где лежат подпапки классов. Параметр image_size задаёт высоту и ширину после загрузки. Партия собирается по параметру batch_size. Ниже тот же приём для короткого пути photos/train:

import os
import tensorflow as tf

root = "photos/train"
for cls in ("apples", "pears"):
	folder = os.path.join(root, cls)
	os.makedirs(folder, exist_ok=True)
	img = tf.cast(tf.ones((6, 6, 3)) * 200, tf.uint8)
	tf.io.write_file(
		os.path.join(folder, "0.png"),
		tf.io.encode_png(img),
	)

ds = tf.keras.utils.image_dataset_from_directory(
	root,
	image_size=(16, 16),
	batch_size=4,
)
for images, labels in ds.take(1):
	print(images.shape)
	# выведет (2, 16, 16, 3)

Метки классов - целые номера: порядок подпапок задаётся алфавитом имён. Если файлов больше, чем места в одной партии, набор делится на несколько шагов:

import os
import tensorflow as tf

base = "mini_album"
for cls in ("cats", "dogs"):
	folder = os.path.join(base, cls)
	os.makedirs(folder, exist_ok=True)
	for i in range(2):
		img = tf.cast(tf.ones((8, 8, 3)) * 128, tf.uint8)
		path = os.path.join(folder, f"{i}.png")
		tf.io.write_file(path, tf.io.encode_png(img))

ds = tf.keras.utils.image_dataset_from_directory(
	base,
	image_size=(8, 8),
	batch_size=2,
	shuffle=False,
)
print(len(list(ds)))
# выведет 2

Создайте каталог samples с подпапками a и b, положите в каждую по одному маленькому файлу. Загрузите набор с размером 12×12 и выведите форму первой партии картинок.

Подготовьте три класса в папке icons и постройте набор с партией из восьми элементов. Выведите тип меток первой партии.

Разложите файлы по двум подкаталогам old и new. Получите набор без перемешивания и выведите метки первой партии.