独自の画像からTensorflowのCNNの学習に利用可能なデータセットを作成する

はじめに

tensorflowのCNNで，独自のデータセットを用いる時に，プログラムに入れるデータを作るプログラムを毎回書くのが面倒だったので，データセットを作成するための関数(image_dataset.py)を作成した．

またただデータを分類するだけでなく，画像のリサイズや，ランダム位置でのクリップをオプションで同時に行えるようにした．

データ構造は以下の構造に従うとし，クラスごとにフォルダ分けされているとする．

./dataset
|
|- dataA
|   |- image1a.jpg
|   |- image2a.jpg
|   |- image3a.jpg
|   |-　 　　　　　・
|   |-　 　　　　　・
|
|- dataB
|   |- image1b.jpg
|   |- image2b.jpg
|   |-　 　　　　　・
|   |-　 　　　　　・
|
・
・

ソースコードの全体像

image_dataset.py

import sys , os
import random

import numpy as np
import cv2


# 画像をランダムの位置で切り抜くプログラム
def random_clip(img , clip_size , num = 1):
    clip_images = []
    height, width = img.shape[:2]

    # 画像をclip_sizeサイズごとにnum回切り抜く
    for y in range( num ):
        rand_y = random.randint(0,height - clip_size)
        rand_x = random.randint(0,width - clip_size)        
        clip_img = img[ rand_y : rand_y + clip_size, rand_x : rand_x + clip_size]
        clip_img = clip_img.flatten().astype(np.float32)/255.0
        clip_images.append(clip_img)

    return clip_images


# データセットから画像とラベルをランダムに取得
def random_sampling( images , labels , train_num , test_num = 0  ):
    image_train_batch = []
    label_train_batch = []

    #乱数を発生させ，リストを並び替える．
    random_seq = list(range(len(images)))
    random.shuffle(random_seq)

    # バッチサイズ分画像を選択
    image_train_batch = images[ :train_num ]
    label_train_batch = labels[ :train_num ]

    if test_num == 0: # 検証用データの指定がないとき
        return image_train_batch , label_train_batch
    else:
        image_test_batch = images[ train_num : train_num + test_num ]
        label_test_batch = labels[ train_num : train_num + test_num ]
        return image_train_batch , label_train_batch , image_test_batch , label_test_batch


# フォルダーの画像をランダム位置でクリップした後にリサイズして読み込む
def make( folder_name , img_size = 0 , clip_num = 0 , clip_size = 0 ,train_num = 0 , test_num = 0 ):
    train_image = []
    test_image = []
    train_label = []
    test_label= []

    # フォルダ内のディレクトリの読み込み
    classes = os.listdir( folder_name )

    for i, d in enumerate(classes):
        files = os.listdir( folder_name + '/' + d  )

        tmp_image = []
        tmp_label = []
        for f in files:
            # 1枚の画像に対する処理
            if not 'jpg' in f:# jpg以外のファイルは無視
                continue

            # 画像読み込み
            img = cv2.imread( folder_name+ '/' + d + '/' + f)
            # one_hot_vectorを作りラベルとして追加
            label = np.zeros(len(classes))
            label[i] = 1

            # リサイズをする処理
            if img_size != 0:
                img = cv2.resize( img , (img_size , img_size ))
                img = img.flatten().astype(np.float32)/255.0
                tmp_image.append(img)                
                tmp_label.append(label)
            elif clip_size != 0 and clip_num != 0:
                img = random_clip( img , clip_size , clip_num)
                tmp_image.extend( img )                    
                for j in range(clip_num):
                    tmp_label.append(label)
            else:
                img = img.flatten().astype(np.float32)/255.0
                tmp_image.append(img)
                tmp_label.append(label)


        # データセットのサイズが指定されているときはその枚数分ランダムに抽出する
        if train_num == 0 :
            train_image.extend( tmp_image )
            train_label.extend( tmp_label )
        #　テスト画像の指定がないとき
        elif test_num == 0 :
            sampled_image , sampled_label = random_sampling( tmp_image , tmp_label , train_num )
            train_image.extend( sampled_image )
            train_label.extend( sampled_label )
        else :
            sampled_train_image , sampled_train_label , sampled_test_image , sampled_test_label = random_sampling( tmp_image , tmp_label , train_num , test_num )
            train_image.extend( sampled_train_image )
            train_label.extend( sampled_train_label )
            test_image.extend( sampled_test_image )
            test_label.extend( sampled_test_label )

        print(d + 'read complete ,' + str(len(train_label)) + ' pictures exit')

    # numpy配列に変換
    train_image = np.asarray( train_image )
    train_label = np.asarray( train_label )

    if test_num != 0: #testデータセットがあるときは返り値が変わる
        test_image = np.asarray( test_image )
        test_label = np.asarray( test_label )
        return train_image , train_label , test_image , test_label

    return train_image , train_label

使い方

通常のファイルと同様にimportする．

python

import image_dataset
train_image ,train_label, test_image ,test_label = image_dataset.make('データセットのディレクトリ'  ,train_num=10000 , test_num =1000)

train_imageとtrain_labelに，train_num枚の画像とそれに対応するラベルが格納され，test_imageとtest_labelにtest_num枚分の画像が格納される．

画像のリサイズを同時に行いたい時

image_dataset.make()の引数にimg_size='サイズ'を指定することで，そのサイズの一片を持つ正方領域でリサイズされる．

例：画像を32×32でリサイズしたものを，学習用セット10000枚，テストセット1000枚ずつ用意する．

python

train_image ,train_label, test_image ,test_label = image_dataset.make('データセットのディレクトリ'  ,img_size = 32, train_num=10000 , test_num =1000)

ランダム位置の切り抜きを同時に行う時

同様に引数にclip_size='サイズ'，clip_num = 'クリップする回数'を指定することで，1枚の画像あたり，clip_num回，clip_size×clip_sizeの矩形領域でランダムに切り抜かれる．

例：1枚の画像を32×32の領域で，ランダムに10回切り抜いたものを，学習用セット10000枚，テストセット1000枚ずつ用意する．

python

train_image ,train_label, test_image ,test_label = image_dataset.make('データセットのディレクトリ'  ,clip_size = 32, clip_num = 10 , train_num=10000 , test_num =1000)

参考

TensorFlowとCNNで、自作データセットを画像分類する
TensorFlowでキルミーアイコン686枚によるキルミー的アニメ絵分類