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

はじめに

• 学習には大量の訓練データが必要になります。
• 目安として分類の場合、1ラベル訓練データ:10,000枚。テストデータ:3,000枚くらいと耳にしたりします。
• 畳み込みを行う際は、よく画像が正方形になってます。しかしながら、この正方形画像を用意するのは大変です。
• 今回は無理やり隙間を埋めて正方形画像にリサイズすることにしました。
• ただ、この隙間が学習結果にどのように影響を及ぼすのかはまだ検証中のため、保証できません。

プログラム

resize.py

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

import argparse
import os
import cv2
import numpy as np

def resizeImages(input_dir, output_dir):

    files = os.listdir(input_dir)

    for file in files:

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

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

        img = cv2.imread(
            os.path.join(input_dir, file),
            cv2.IMREAD_COLOR
        )

        tmp = img[:, :]
        height, width = img.shape[:2]

        if (height > width):
            size = height
            limit = width
        else:
            size = width
            limit = height

        start = int((size - limit) / 2)
        fin = int((size + limit) / 2)

        resized_img = cv2.resize(np.zeros((1, 1, 3), np.uint8), (size, size))

        if (size == height):
            resized_img[:, start:fin] = tmp
        else:
            resized_img[start:fin, :] = tmp

        cv2.imwrite(
            os.path.join(output_dir, file),
            resized_img
        )

def main():

    parser = argparse.ArgumentParser()
    parser.add_argument('--input_dir', default='original')
    parser.add_argument('--output_dir', default='resized')
    args = parser.parse_args()

    resizeImages(args.input_dir, args.output_dir)

if __name__ == '__main__':

    main()

実行結果