ImageDataGenerator flow_from_directoryを使い、名前付きでクラスごとのデータを読み込む

kerasの ImageDataGenerator を使って画像を読み込み、kaggleの画像分類問題をやっていたのですが、 validationデータで良い正答率が出るにもかかわらず、testデータにするとうまく分類できない状況に陥りました。

原因は flow_from_directory で classes を指定していなかったことでした。

現象

前提として、データは以下のようなディレクトリ構造で配置しています。

data/
    train/
        dogs/
            dog001.jpg
            dog002.jpg
            ...
        cats/
            cat001.jpg
            cat002.jpg
            ...
    validation/
        dogs/
            dog001.jpg
            dog002.jpg
            ...
        cats/
            cat001.jpg
            cat002.jpg
            ...

続いてコードです。

    train_generator = train_datagen.flow_from_directory(
        train_images, # 上記 train を指定
        target_size=(image_height, image_width),
        batch_size=batch_size,
        class_mode='categorical')

    validation_generator = test_datagen.flow_from_directory(
        test_images, # 上記 validation/ を指定
        target_size=(image_height, image_width),
        batch_size=1,
        class_mode='categorical')

flow_from_directory は class_mode を categorical にしておけば上記のディレクトリ構造からいい感じに分類をやってくれるそうです。 ということで、分類を test/ や validation/ 以下のディレクトリ名でやってくれると思ったのですが、そうではありませんでした。 このまま実行するとディレクトリ名に関係なく0からディレクトリ数までの数字を割り振って分類されます。 つまりdogs, catsではなく0, 1のように分類されるということです。

解決策

この問題を解決するためには classes を指定してやる必要があります。

    classes = ['dogs', 'cats']

    train_generator = train_datagen.flow_from_directory(
        train_images, # 上記 train を指定
        target_size=(image_height, image_width),
        batch_size=batch_size,
        classes=classes,
        class_mode='categorical')

    validation_generator = test_datagen.flow_from_directory(
        test_images, # 上記 validation/ を指定
        target_size=(image_height, image_width),
        batch_size=1,
        classes=classes,
        class_mode='categorical')