畳み込みニューラルネットワーク(CNN)の訓練データを作成するツールを作ってみた - ③データリストの作成

はじめに

• 前回と前々回に続いて、今回は訓練データとテストデータのファイルリストを作成を行います。
• このファイルリストを元に実際は学習プログラムを実行していく想定で作成しました。

プログラム

filelist.py

# !/usr/local/bin/python3
# !-*- coding: utf-8 -*-

import os
import random
import numpy as np

def createFilelist():

    labels = {"gyoza":0, "pizza":1, "sushi":2}

    dir_path = os.path.dirname(os.path.abspath(__file__))

    for label_name, label_key in labels.items():

        files = os.listdir(label_name)
        random.shuffle(files)
        files = np.array(files)

        train, test = np.split(files, [int(files.size * 0.7)])

        print(
            'ラベル名: ' + label_name,
            '訓練データ数: ' + str(len(train)),
            'テストデータ数: ' + str(len(test)),
        )

        for file in train:

            name, ext = os.path.splitext(file)

            if ext != '.jpg':
                print('[' + file + ']: 不正なファイルが含まれています。')

            with open('train.txt', 'a') as f:
                text = '{0}/{1}/{2}\t{3}\n'.format(dir_path, label_name, file, labels[label_name])
                f.write(text)

        for file in test:

            name, ext = os.path.splitext(file)

            if ext != '.jpg':
                print('[' + file + ']: 不正なファイルが含まれています。')

            with open('test.txt', 'a') as f:
                text = '{0}/{1}/{2}\t{3}\n'.format(dir_path, label_name, file, labels[label_name])
                f.write(text)

def main():

    createFilelist()

if __name__ == '__main__':

    main()

実行結果

{プロジェクト絶対パス}/gyoza/2827573.jpg	0
{プロジェクト絶対パス}/gyoza/2628927.jpg	0