tensorflow 이미지 데이터셋 만들기

1. 이미지 수집

a1.jpg

b1.jpg

위와 같이 그림판으로 'a' 를 그린 이미지 10개, 'b' 를 그린 이미지 5개를 각각 다른 이름으로 한 디렉토리에 저장한다.

./

 

a1.jpg
a2.jpg
a3.jpg
a4.jpg
a5.jpg
b1.jpg
b2.jpg
b3.jpg
b4.jpg
b5.jpg

2. 코드 작성

./main.py

import tensorflow as tf
import glob
import numpy as np

# 이미지 경로 불러오기 및 정렬
image_paths = sorted(glob.glob('./*.jpg'))

# 이미지 처리 함수
def load_and_preprocess_image(path):
    image = tf.io.read_file(path)
    image = tf.image.decode_jpeg(image, channels=1)  # 흑백 이미지
    image = tf.image.resize(image, [28, 28])  # 크기 조절
    image /= 255.0  # 정규화
    return image

# 이미지 데이터셋 생성
path_ds = tf.data.Dataset.from_tensor_slices(image_paths)
image_ds = path_ds.map(load_and_preprocess_image, num_parallel_calls=tf.data.AUTOTUNE)

# 라벨 데이터셋 생성
labels = [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]
label_ds = tf.data.Dataset.from_tensor_slices(labels)

# 이미지와 라벨 결합
image_label_ds = tf.data.Dataset.zip((image_ds, label_ds))

# 데이터셋 셔플 및 분할
DATASET_SIZE = len(image_paths)
TRAIN_SIZE = int(0.8 * DATASET_SIZE)
SHUFFLE_BUFFER_SIZE = 1000

# 데이터셋 셔플
shuffled_ds = image_label_ds.shuffle(SHUFFLE_BUFFER_SIZE)

# 훈련 데이터셋
train_ds = shuffled_ds.take(TRAIN_SIZE)

# 테스트 데이터셋
test_ds = shuffled_ds.skip(TRAIN_SIZE)

# NumPy 배열로 변환
train_images = np.array([image.numpy() for image, _ in train_ds])
train_labels = np.array([label.numpy() for _, label in train_ds])
test_images = np.array([image.numpy() for image, _ in test_ds])
test_labels = np.array([label.numpy() for _, label in test_ds])

# npz 파일로 저장
np.savez_compressed('datasets.npz', 
                    train_images=train_images, train_labels=train_labels,
                    test_images=test_images, test_labels=test_labels)

3. 실행

python3 main.py

datasets.npz 파일이 생성됐다면 성공!