chainerのLabeledImageDatasetで読み込める画像リストのテキストを作成する

chainerのtrainerには、LabeledImageDatasetという学習データを読み込むフォーマットが用意されている。このクラスには画像のパスとラベルが書かれたテキストファイルのパスを入力できる。これにより、バッチごとに逐次画像を読み込んでくれるのでメモリの節約になるし、自分で画像を読み込むコードを書かなくていい。

今回はそのテキストを作成するための手順を示す。

chainer.datasets.LabeledImageDataset
https://docs.chainer.org/en/stable/reference/generated/chainer.datasets.LabeledImageDataset.html

LabeledImageDatasetで読み込めるテキスト

画像のパスとそのラベルの値をスペースで分けておく必要がある。

dataset.txt

a.jpg 0
b.jpg 1

動作環境

• Ubuntu 16.04 LTS
• chainer 4.0.0

解説

以下のようなファイル構成で画像が格納されているとする。

$ tree image
image/
├── 00
│   ├── a.JPG
│   ├── b.JPG
│   ├── c.JPG
│   ├── d.JPG
│   ├── e.JPG
│   ├── f.JPG
│   ├── g.JPG
│   └── h.JPG
├── 01
│   ├── a.JPG
│   ├── b.JPG
│   ├── c.JPG
│   ├── d.JPG
│   ├── e.JPG
│   ├── f.JPG
│   ├── g.JPG
│   └── h.JPG
├── 02
│   ├── a.JPG
│   ├── b.JPG
│   ├── c.JPG
│   ├── d.JPG
│   ├── e.JPG
│   ├── f.JPG
│   ├── g.JPG
│   └── h.JPG
├── 03
│   ├── a.JPG
│   ├── b.JPG
│   ├── c.JPG
│   ├── d.JPG
│   ├── e.JPG
│   ├── f.JPG
│   ├── g.JPG
│   └── h.JPG
└── 04
    ├── a.JPG
    ├── b.JPG
    ├── c.JPG
    ├── d.JPG
    ├── e.JPG
    ├── f.JPG
    ├── g.JPG
    └── h.JPG

以下のように実行すると、train_XXX.txtとtest_XXX.txtが生成される。

$ python3 ./write_dataset.py image

catコマンドで確認すると以下のようにファイルが作成されていることが確認できる。今回はimage直下のディレクトリが00から04と5個あるのでtest_005.txtというファイル名になっている。

$ cat result/test_005.txt 
./image/00/e.JPG 0
./image/02/g.JPG 2
./image/01/c.JPG 1
./image/04/c.JPG 4

学習時は以下のようにし、その後にiteratorに格納して使用する。

train.py

train = LabeledImageDataset('train_005.txt')
test = LabeledImageDataset('test_005.txt')

ソースコード

write_dataset.py

import os
import argparse
import numpy as np
from glob import glob

import imgfunc as IMG


def command():
    parser = argparse.ArgumentParser(description=help)
    parser.add_argument('img_root_path',
                        help='テキストデータを作成したいデータセットのルートパス')
    parser.add_argument('--train_per_all', '-t', type=float, default=0.9,
                        help='画像数に対する学習用画像の割合 [default: 0.9]')
    parser.add_argument('-o', '--out_path', default='./result/',
                        help='データの保存先 (default: ./result/)')
    return parser.parse_args()


def writeTXT(folder, name, data):
    """
    テキストファイルを書き出す
    [in] folder: テキストを保存するフォルダ
    [in] name:   テキストの名前
    [in] data:   保存するデータ

    dataは[(path1, val1), (path2, val2), ... , (pathN, valN)]の形式であること
    pathN: N番目の画像パス
    valN:  N番目の画像の分類番号
    """

    with open(os.path.join(folder, name), 'w') as f:
        [f.write('./' + i + ' ' + j + '\n') for i, j in data]


def str2int(in_str):
    val = 0
    try:
        val = int(in_str)
    except:
        print('ERROR:', in_str)
        val = -1

    return val


def main(args):
    # 画像データを探索し、画像データのパスと、サブディレクトリの値を格納する
    search = glob(os.path.join(args.img_root_path, '**'), recursive=True)
    data = [(img, str2int(img.split('/')[-2])) for img in search
            if IMG.isImgPath(img, True)]
    # ラベルの数を数える
    label_num = len(np.unique(np.array([i for _, i in data])))
    print('label num: ', label_num)
    # 取得したデータをランダムに学習用とテスト用に分類する
    data_arr = np.array(data)
    data_len = len(data_arr)
    shuffle = np.random.permutation(range(data_len))
    train_size = int(data_len * args.train_per_all)
    train = data_arr[shuffle[:train_size]]
    test = data_arr[shuffle[train_size:]]
    # chainer.datasets.LabeledImageDataset形式で出力する
    writeTXT(args.out_path, 'train_' + str(label_num).zfill(3) + '.txt', train)
    writeTXT(args.out_path, 'test_' + str(label_num).zfill(3) + '.txt', test)


if __name__ == '__main__':
    args = command()
    main(args)

imgfunc.py

import cv2

def isImgPath(name, silent=False):
    """
    入力されたパスが画像か判定する
    [in]  name:   画像か判定したいパス
    [in]  silent: cv2.imread失敗時にエラーを表示させない場合はTrue
    [out] 画像ならTrue
    """

    if not type(name) is str:
        return False

    # cv2.imreadしてNoneが返ってきたら画像でないとする
    if cv2.imread(name) is not None:
        return True
    else:
        if not silent:
            print('[{0}] is not Image'.format(name))

        return False

以上。