スマブラーが動画から数値を認識させてみた(後編)

前回の内容

スマブラーが動画から数値を認識させてみた(前編)

ダメージの認識には背景の情報量を落とす必要がありそうですっというところまで書きました。

グレースケール化

スマブラではダメージが蓄積していくと表記が赤色になっていくので、青色を除外した状態でグレースケール化するやり方と、そのままグレースケール化するやり方を試しました。
どちらもそこそこキレイに抽出出来ますが、背景の色によって(黒だったり赤だったり)は上手くいかない可能性も考えて２つをマージしてます。
※まだ検証不十分なので今後最適化する予定ですが、今のところはうまくいってます

sample.py

    def get_image_for_model(self, image: np.ndarray) -> np.ndarray:
        image_cp = image.copy()

        # pick up number(by remove blue)
        grn = image_cp[:, :, 1]
        red = image_cp[:, :, 2]
        red_grn = cv2.addWeighted(red, 0.5, grn, 0.5, 0)
        th_red_grn = cv2.adaptiveThreshold(red_grn, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY, 51, 10)

        # another pick up number(by gray scale threshold)
        grayed = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
        ret, th_gr = cv2.threshold(grayed, 10, 255, cv2.THRESH_BINARY)

        # merge
        merged = np.minimum(th_red_grn, th_gr)

        return merged

output

Tesseractを使えるか

答えはNoでした。完璧にゴミを消すことは出来ず、Tesseractではどうしてもゴミを文字として認識してしまうことがあるため、機械学習させることにしました。

機械学習

今回はサポートベクターマシンを使うことにしました。
ロジスティック回帰よりもマージンをとって分類してくれるので精度が上がるかなと思ったので採用しています。

機械学習させるにあたり、トレーニングデータの用意に少し悩みました。
結論としては下記のようなフォルダ構成を用意しました。

まずはInputフォルダに最初は少量の画像データを入れて、そのデータで学習させてモデルを作ります。そして作ったモデルで分類を行いOutputフォルダに結果を格納します。その後、Outputフォルダの中を確認して、不正解な分類をした画像を少しずつInputフォルダに追加していきました。

フォルダ構成.

各数字に対応した画像データを格納 ※数字ではない画像データは10に格納
Input　　#トレーニングデータフォルダ
　├── 0
　├── 1
　├── 2
　├── 3
　├── 4
　├── 5
　├── 6
　├── 7
　├── 8
　├── 9
　└── 10


Output　　#分類結果フォルダ
　├── 0
　├── 1
　├── 2
　├── 3
　├── 4
　├── 5
　├── 6
　├── 7
　├── 8
　├── 9
　└── 10

make_model.py

import cv2
import numpy as np
import os
import pickle
from sklearn.svm import SVC
from sklearn.metrics import accuracy_score
from sklearn.model_selection import train_test_split
from pathlib import Path

#########################
# SETTING
#########################
train_data_dir = 'E:/sample/Input'             # training data directory
output_pickle = 'testmodel.pickle'             # output pickle file name
#########################

X = []
y = []

print('start reading training data :)')
p_label_dir = Path(train_data_dir)
it_dir = p_label_dir.iterdir()

# read image files from input directory
for dir in it_dir:
    if dir.is_dir():
        p_filedir = Path(dir)
        it_file = p_filedir.glob('*.png')
    else:
        continue

    for file in it_file:
        file = str(file)
        image = cv2.imread(file, 0)
        X.append(image.reshape(-1,))
        y.append(int(dir.name))

X = np.array(X)
y = np.array(y)

print('read complete :D')

print('start training :O')
X_train, X_test, y_train, y_test = train_test_split(X, y)

model = SVC(kernel='linear', random_state=None)
model.fit(X_train, y_train)
print('complete training :P')

print('culculate accuracy score...')
pred_train = model.predict(X_train)
accuracy_train = accuracy_score(y_train, pred_train)
print('accuracy score at training data set： %.2f' % accuracy_train)

print('culculate accuracy score...')
pred_test = model.predict(X_test)
accuracy_test = accuracy_score(y_test, pred_test)
print('accuracy score at test data set： %.2f' % accuracy_train)

filepath = os.path.split(__file__)[0]
uri_pickle = os.path.join(filepath, output_pickle)

with open(uri_pickle, mode='wb') as fp:
    pickle.dump(model, fp)

classify_numbers.py

import cv2
import os
import recognition
import pickle

#########################
# SETTING
#########################
test_data_dir = 'E:/sample/testdata'           # test data directory
picklefile = 'testmodel.pickle'                # pickle file of model data
output_dir = 'E:/sample/Output'                # output directory
progress_step = 1000                           # step count for debug log
#########################

filepath = os.path.split(__file__)[0]
uri_pickle = os.path.join(filepath, picklefile)

with open(uri_pickle, mode='rb') as fp:
    model = pickle.load(fp)

# get file names from test data directory
file_list = os.listdir(test_data_dir)

for i, file_name in enumerate(file_list):
    if(i % progress_step == 0):
        print('Read progress rate {}/{}'.format(i, len(file_list)))

    abs_name = os.path.join(test_data_dir, file_name)
    image_org = cv2.imread(abs_name, 0)

    number = str(model.predict(image_org.reshape(1, -1))[0])

    output_sub_dir = os.path.join(output_dir, number)
    output_filename = os.path.join(output_sub_dir, '{}.png'.format(i))

    cv2.imwrite(output_filename, image_org)

反省

ダメージ値を認識させ、グラフ化は出来ましたが、イマイチ有用ではありませんでした。
抽出したデータをどのように使うかが大事だと感じました。(現在いろいろ試行錯誤中です。。。)

また、Qiitaのような技術系の投稿も初めてなこともあり、今見直してみるとイマイチ誰向けなのかピンと来ない投稿になってしまいました。現在進行形なのですが、以下のようなデスクトップアプリをEelを使って作成中です。

Eelの記事がとても少ないので、私自身苦労していますが、とても便利なフレームワークなので、次回はEel関連の投稿をしてみようと思います。