AIについて勉強を始めた際、最初にたどり着くのは、MNISTのデータセットと言う人も多いと思います。MNISTのデータセットは、画像のサイズが小さいため、非常に軽量で、GPUの入っていないPCでも学習を行うことのできるデータセットの一つです。

また、MNISTのデータセットはPythonのAI用のライブラリからも簡単に取得できるため、気軽に使うことができます。一方で、AI初心者はPythonの初心者であることも多く、MNISTのデータセットで動かした後に、別のデータを使って学習をさせようとした際に、苦労する人も多いと思います。

そこで今回は、そのまま使えるMNISTのデータセットをあえて一度フォルダにダウンロードを行い、再度データをロードする方法を学ぶことで、MNISTのデータセット以外のデータでも自力でAI開発を行う方法を習得します。

まずは、今回は、MNISTのダウンロードにkerasを使用し、画像の保存にOpenCVを使用します。

まずは、MNSITのデータセットをダウンロードし、フォルダに保存するソースコードです。

from keras.datasets import mnist
import cv2
import os


def save_images(images, labels, save_dir, prefix):
    """
    画像[images]を指定のフォルダ[save_sir]/[label]/に保存する。
    ファイル名は、[prefix]_[no].png
    :param images: イメージデータ
    :param labels: ラベルデータ
    :param save_dir: 保存ディレクトリ
    :param prefix: 接頭語
    """
    print('Start sabe images:', prefix)
    # ラベルごとに保存先のディレクトリを作成
    # ラベル名の一覧を取得([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])
    # numpy.ndarrayを配列(list)に変換
    labels_list = labels.tolist()
    # listから重複しているラベル名を排除する。
    class_list = list(set(labels_list))
    print(class_list)

    for class_name in class_list:
        # ラベルごとのディレクトリを作成する
        label_save_dir = os.path.join(save_dir, str(class_name))
        os.makedirs(label_save_dir, exist_ok=True)

    no = 0
    for image, label in zip(images, labels):
        # 接頭語_番号.pngの命名で自分のラベルに保存
        file_name = prefix + '_' + str(no).zfill(5) + '.png'
        file_path = os.path.join(save_dir, str(label))
        file_path = os.path.join(file_path, file_name)

        cv2.imwrite(file_path, image)
        no += 1


def mnist_download_and_save(save_dir):
    """
    MNISTのデータをダウンロードし、指定のフォルダ[save_dir]に保存する。
    :param save_dir: データの保存先ディレクトリ
    """
    (x_train, y_train), (x_test, y_test) = mnist.load_data()
    print(x_train.shape, x_test.shape)

    # 学習用データをディレクトリに保存する
    train_dir = os.path.join(save_dir, 'train')
    save_images(x_train, y_train, train_dir, 'train')

    # テスト用データをディレクトリに保存する
    test_dir = os.path.join(save_dir, 'test')
    save_images(x_test, y_test, test_dir, 'test')


if __name__ == '__main__':
    mnist_download_and_save('./mnist_data')

また、保存したデータを読み取る方法は、以下の通りとなります。

import cv2
import os
import numpy as np


def image_data_load_for_dir(image_dir, class_list):
    """
    ディレクトリから画像データとラベルデータを取得する
    :param image_dir: 画像データディレクトリ
    :param class_list: クラス名のリスト
    :return: x_data, y_data
    """

    print('Start image_data_load_for_dir:', image_dir)
    x_data = []
    y_data = []
    for class_name in class_list:
        # ラベルごとのディレクトリ名を取得
        target_image_dir = os.path.join(image_dir, str(class_name))
        # os.listdirでディレクトリ内のファイル名一覧を取得する
        for file_name in os.listdir(target_image_dir):
            file_path = os.path.join(target_image_dir, file_name)
            # 画像の取得
            image = cv2.imread(file_path, cv2.IMREAD_GRAYSCALE)
            x_data.append(image)
            y_data.append(class_name)

    x_data = np.array(x_data)
    y_data = np.array(y_data)
    return x_data, y_data


def mnist_data_load(root_dir, class_list):
    """
    MNISTのデータをロードします
    :param root_dir: 学習データおよびテストデータのルートディレクトリ
    :param class_list: クラス名のリスト
    """
    # 学習用データをディレクトリに保存する
    train_dir = os.path.join(root_dir, 'train')
    x_train, y_train = image_data_load_for_dir(train_dir, class_list)
    print('学習データ', x_train.shape, y_train.shape)

    # テスト用データをディレクトリに保存する
    test_dir = os.path.join(root_dir, 'test')
    x_test, y_test = image_data_load_for_dir(test_dir, class_list)
    print('テストデータ', x_test.shape, y_test.shape)


if __name__ == '__main__':
    mnist_class_list = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
    mnist_data_load('./mnist_data', mnist_class_list)

終わりに

自分が学習に使っているデータセットがどのようなデータであるかを確認することは非常に重要な作業です。必ず確認し、データの特徴等を確認しましょう。

また、今回使用した画像の読み取りなどがわからない人は、以下の記事も確認し何を行なっているか確認しましょう。

参考

コメントを残す

メールアドレスが公開されることはありません。 * が付いている欄は必須項目です