茨城エンジニアのPython開発日記

茨城のITベンチャー企業ではたらく2年目エンジニア。Pythonで色々なものを作成中。

フォルダ構造の変更で識別カテゴリを変更, アノテーションまでできる画像識別プログラムを作ってみた


ブログから記事を見つけたい場合はこちら

ブログ地図 - 茨城エンジニアのPython開発日記


あけましておめでとうございます。松原です。
年末は初めて格闘技みてましたけど、結構おもしろかったです。
来年もみてみようかな。


さて、今回はフォルダ構造から識別カテゴリを取得してくれるコードを作りました。
フォルダ自体が正解値を与えてくれるので、アノテーションもフォルダに振り分けた時点で完了しています。
f:id:tottorisnow33:20210105003336p:plain
f:id:tottorisnow33:20210105003356p:plain

こんな感じに学習用画像と評価用画像をいれておけば、中のファイルを全て使って学習してくれます。
カテゴリ数も可変なので、好きなようにフォルダ増やしてカテゴリをいじってください。
ちなみにソースコードはこの階層にいれておきます。

f:id:tottorisnow33:20210105003606p:plain

で、ソースコードの役割はこんな感じ。

load.py: パーツの関数置き場。主に学習用画像読み込みに利用。
train.py: 学習実行用コード。学習したいときはこのファイルを実行。学習結果h5ファイルとして出力される。
demo.py: 学習してできたh5ファイルを使って推論を実施。demo_imgフォルダ直下のjpg全てに推論をかける。

ソースコードの中身はこんな感じ。

1.load.py

#返り値として訓練用データセットが返る関数

#訓練データ
#x_train[画像枚数][高さ][幅][RGB]            入力
#y_train[画像枚数][クラス数]                 出力

#テストデータ
#x_test
#y_test
from PIL import Image
import numpy as np
import glob
import os
import cv2

########################################################################
###メイン処理, 要求された画像数のx_train, y_trainm, x_test, y_testを返す###
########################################################################
def Get_Training_and_Test_Image():

    #データセット内のデータセット群を取得
    dataset_folders_path = Get_dataset_name()

    #カテゴリ名取得
    dataset_cate_name = Get_dataset_name()

    #画像データサイズを決定(このサイズにリサイズする)
    img_size=(128, 128)

    #データから入力行列, 出力行列を作成(訓練用)
    x_train, y_train = Get_xy_data(dataset_folders_path, dataset_cate_name, img_size, 0)

    #データから入力行列, 出力行列を作成(評価用)
    x_test, y_test = Get_xy_data(dataset_folders_path, dataset_cate_name, img_size, 1)

    print(type(x_train), x_train.shape)
    print(type(y_train), y_train.shape)
    print(type(x_test), x_test.shape)
    print(type(y_test), y_test.shape)

    return x_train, y_train, x_test, y_test, dataset_cate_name


#################################
#データセットのサブフォルダ群を取得#
#################################
def Get_dataset_name():
    
    #カレントディレクトリ直下のdatasetフォルダを探索
    dataset_path = os.getcwd() + "\dataset"

    #サブフォルダーのパスを取得
    folders = os.listdir(dataset_path)
    folders_name = [f for f in folders if os.path.isdir(os.path.join(dataset_path, f))]

    return folders_name

#########################
#x_data, y_dataを返す関数#
#########################
def Get_xy_data(dataset_folders_path, dataset_cate_name, img_size, mode):

    #行列用意
    x_mats, y_mats = [], []

    #x_dataのパス用意
    x_data_paths = ["temp" for i in range(len(dataset_folders_path))]

    if(mode == 0):
        #訓練データの行列を取得するモード
        x_data_paths = ["dataset/" + dataset_folders_path[i] + "/train" for i in range(len(dataset_folders_path))]
    elif(mode == 1):
        #評価データの行列を取得するモード
        x_data_paths = ["dataset/" + dataset_folders_path[i] + "/test" for i in range(len(dataset_folders_path))]
    print(x_data_paths)

    #全サブフォルダを処理
    for folder_cnt in range(len(x_data_paths)):

        #フォルダ内のファイル取得
        img_file_names = os.listdir(x_data_paths[folder_cnt])

        #フォルダ内の画像を処理
        for img_file_cnt in range(len(img_file_names)):
            
            #画像の相対パス取得
            img_file_path = x_data_paths[folder_cnt] + "/" + img_file_names[img_file_cnt]

            #画像の前処理
            image = preparation_img(img_file_path, img_size)
           
            #画像をリストに追加
            x_mats.append(image)

            #正解値をリストに追加(サブフォルダの番号がそのままカテゴリ番号)
            y_mat = []
            for i in range(len(x_data_paths)):
                #フォルダカウントの要素に1をいれる
                if(i == folder_cnt):
                    y_mat.append(1)
                else:
                    y_mat.append(0)
            #今回の画像の正解値をy_matsに格納
            y_mats.append(y_mat)
            
        
    x_mats = np.asarray(x_mats, dtype=np.float32)
    y_mats = np.asarray(y_mats, dtype=np.uint8)

    #np.savetxt(img_file_names[img_file_cnt] + "_x.csv" ,x_mats[0,:,:,0], delimiter=',')

    return x_mats, y_mats

#########################
########画像前準備########
#########################
def preparation_img(img_file_path, img_size):

    #ファイルを開く
    image = Image.open(img_file_path)

    #正方形へ変換
    image = crop_to_square(image)

    #リサイズ
    image = image.resize(img_size)

    #RGBA→RGB
    if image.mode == "RGBA":
        image = image.convert("RGB")
    
    #numpy配列にして小数化
    image = np.asarray(image)
    image = image / 255.0

    return image


#############################################
########画像を正方形に変換する処理(借用)########
#############################################
def crop_to_square(image):
    size = min(image.size)
    left, upper = (image.width - size) // 2, (image.height - size) // 2
    right, bottom = (image.width + size) // 2, (image.height + size) // 2
    return image.crop((left, upper, right, bottom))

Get_Training_and_Test_Image()

2.train.py

import tensorflow as tf
from load import Get_Training_and_Test_Image

# グラボのメモリが不足してる時はこれをいれる
#グラボがないときはコメントアウト
physical_devices = tf.config.list_physical_devices('GPU')
if len(physical_devices) > 0:
    for device in physical_devices:
        tf.config.experimental.set_memory_growth(device, True)
        print('{} memory growth: {}'.format(device, tf.config.experimental.get_memory_growth(device)))
else:
    print("Not enough GPU hardware devices available")


#データセット用意
x_train, y_train, x_test, y_test, dataset_cate_name = Get_Training_and_Test_Image()


#epoch数とbatchサイズ設定
set_epoch_num = 10
set_batch_size = 2

######################NN構築(CNN)##############################
input_shape = (128, 128, 3)

# モデル構築
inputs = tf.keras.Input(shape=input_shape, name='input')

x = tf.keras.layers.Conv2D(3,  kernel_size=(3, 3), activation='relu', padding='same')(inputs)
x = tf.keras.layers.MaxPooling2D((2, 2))(x)
x = tf.keras.layers.Conv2D(5,  kernel_size=(3, 3), activation='relu', padding='same')(x)
x = tf.keras.layers.MaxPooling2D((2, 2))(x)
x = tf.keras.layers.Conv2D(6, kernel_size=(3, 3), activation='relu', padding='same')(x)
x = tf.keras.layers.MaxPooling2D((2, 2))(x)
x = tf.keras.layers.Conv2D(3, kernel_size=(3, 3), activation='relu', padding='same')(x)
x = tf.keras.layers.MaxPooling2D((2, 2))(x)
x = tf.keras.layers.Flatten()(x)
x = tf.keras.layers.Dense(len(dataset_cate_name), activation='relu')(x)

outputs = tf.keras.layers.Dense(len(dataset_cate_name), activation='softmax')(x)

#層をセット(作ったパーツの組み合わせ終わったものをここでセッティング)
model = tf.keras.Model(inputs=inputs, outputs=outputs, name="img_recog_model")

model.summary()

######################NN構築##############################

#モデルのコンパイル
model.compile(
#SparseCategoricalCrossentropy
loss=tf.keras.losses.CategoricalCrossentropy(from_logits=False),
optimizer=tf.keras.optimizers.RMSprop(),
metrics=["accuracy"],
)

#学習
model.fit(x_train, y_train, epochs=set_epoch_num, batch_size=set_batch_size)

#評価
score = model.evaluate(x_test, y_test, verbose=2)
print('損失:', score[0])
print('正答率:', score[1])

#モデルの保存
model.save("model_img_recog.h5")

3.demo.py

from PIL import Image
import numpy as np
import glob
import os
import cv2
import tensorflow as tf
from load import Get_dataset_name, preparation_img

img_size = (128, 128)

###########################################
###任意の画像に対して推論→可視化のメイン関数###
###########################################
def demo_main():

   #NNの情報入手
   #同じモデルを読み込んで、重みやオプティマイザーを含むモデル全体を再作成
    model = tf.keras.models.load_model('model_img_recog.h5')
    model.summary()

    gpus = tf.config.experimental.list_physical_devices('GPU')
    if gpus:
        try:
            for gpu in gpus:
                tf.config.experimental.set_memory_growth(gpu, True)
        except RuntimeError as e:
            print(e)
        
    #画像のフォルダパス取得
    img_dir_path ="demo_img"
    img_file_names = os.listdir(img_dir_path)

    #フォルダ内の全ての画像で推論
    for img_cnt in range(len(img_file_names)):

        #画像パス取得
        img_path = img_dir_path + "/" + img_file_names[img_cnt]

        #画像前処理
        input_img = input_img = preparation_img(img_path, img_size)
        input_img = np.asarray(input_img, dtype=np.float32)
        input_img = np.expand_dims(input_img, axis = 0)

        #推論実施
        predict_results = model.predict(input_img)

        #カテゴリ名取得
        folders_name = Get_dataset_name()
        
        #推論結果取得
        predict_result = predict_results[0,:]
        
        #結果出力
        np.set_printoptions(precision=3, suppress=True)
        print("-------------------------------------")
        print(img_file_names[img_cnt])
        print("-------------------------------------")
        for category_cnt in range(len(folders_name)):
            print (folders_name[category_cnt] , ":  " , predict_result[category_cnt]*100, "%")
        
demo_main()


年末年始はこんなものを作ってました。
1日でさっくり終わらせる予定だったんだけど、なぜかどんな画像でも推論結果が同じになる現象が発生。
別データセットで学習しなおすと推論結果は変わるし、入力画像も違う値で入力されることは分かってるのに、同じモデルでやるとどんな画像相手にも同じ推論結果がでてきちゃうんですね。
なんだこれ。

ひとまず年末年始のアウトプットということでここまでの成果を出しときました。
そのうちデバッグしてあげなおします。

※2021/1/5追記
プーリング層入れ忘れてるせいでした。
プーリング層がないと推論結果が同じになるっていうのも不思議だけど……

ちなみにdemo.pyの結果はこんな感じ。

f:id:tottorisnow33:20210105194859p:plain

適当なデータセットもってきてフォルダ振り分けたら結構正解率80%くらいまではさらっといきました。
次回はちょっと精度良くしてご紹介します。