686 views

この記事は最終更新から 395日 が経過しています。

【1】やりたいこと

GIMPなどの画像エディタを使って自筆の数字を描き、これを学習済みパラメータを使って自動認識させてみたい。
作成する画像は、学習に使用したのと同じ 28×28[pixel]の 1[channel]画像（＝モノクロ画像）だ。

【2】やってみる

1) プログラミング

(1) dataset_MNIST.py

１枚の数字画像を読み込み Tensorデータ化する関数 load_one_image を追加した。
変更箇所をハイライト表示する。

from torchvision import datasets, transforms
from PIL import Image

class MyDataset:
    def __init__(self, data_dir='../data'):
        self.transform = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize((0.1307,), (0.3081,))
        ])
        self.data_dir = data_dir

    def get_train_dataset(self):
        return datasets.MNIST(self.data_dir, train=True, download=True, transform=self.transform)

    def get_test_dataset(self):
        return datasets.MNIST(self.data_dir, train=False, download=True, transform=self.transform)

    def load_one_image(self, img_path):
        # 指定された画像ファイルを読み込み、MNIST準拠のTensorに変換して返す。
        # 出力テンソルの shape: [1, 1, 28, 28]
        image = Image.open(img_path).convert('L')         # モノクロに変換
        image = image.resize((28, 28))                    # サイズを28x28に統一
        image = self.transform(image)                     # Tensor化＋正規化
        return image.unsqueeze(0)                         # [1, 1, 28, 28] に拡張

(2) control_test_one_img.py

こちらは新規作成したファイルだ。
control_test.pyを流用し、単一画像を入力＆自動認識するように改造した。

import torch
from net_model_CNN_01 import Net
from trainer import Trainer
from dataset_MNIST import MyDataset
import argparse

#//////////////////////////////////////////////////////////////////////////
def predict_image(img_path, model_path):
    # 環境設定
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    
    # モデルとパラメータの読み込み
    model = Net()
    trainer = Trainer(model, device)
    trainer.load(model_path)

    # 画像読み込みと前処理
    dataset = MyDataset()
    image_tensor = dataset.load_one_image(img_path).to(device)

    # 推論
    model.eval()
    with torch.no_grad():
        output = model(image_tensor)
        prob = torch.softmax(output, dim=1)
        predicted_label = torch.argmax(prob).item()
        prob_array = prob.cpu().numpy()[0]
        formatted = [f"{p:.1f}" for p in prob_array]
        print(f"予測された数字は: {predicted_label}")
        print(f"各クラスの確率: {formatted}")

#//////////////////////////////////////////////////////////////////////////
def main():
    parser = argparse.ArgumentParser(description="Predict a digit from a single image")
    parser.add_argument('-i', type=str, required=True, help='Path to input image file')
    parser.add_argument('-p', type=str, required=True, help='Path to trained model file')
    args = parser.parse_args()
    predict_image(args.i, args.p)

#//////////////////////////////////////////////////////////////////////////
if __name__ == "__main__":
    main()

2) 画像作成

GIMPで 28×28[pixel]のモノクロ画像を作った。
数字は AIが悩みそうなもの、AIが楽々わかりそうなもの、いろいろな数字画像を混ぜてみた。

3) 実行結果

まずは 10 epochs学習する。
　→ 正解率 99.05%のモデルが出来た。

$ python control_train.py -p AAA.pth
[1/10] Epoch complete
[2/10] Epoch complete
[3/10] Epoch complete
[4/10] Epoch complete
[5/10] Epoch complete
[6/10] Epoch complete
[7/10] Epoch complete
[8/10] Epoch complete
[9/10] Epoch complete
[10/10] Epoch complete
Test Accuracy: 99.05%
Model saved to AAA.pth

このモデルを使って、自筆数字画像を自動判読してみる。

$ python control_test_one_img.py -p AAA.pth -i ./Img/t_001_2.png
Model loaded from AAA.pth
予測された数字は: 2
各クラスの確率: ['0.0', '0.0', '1.0', '0.0', '0.0', '0.0', '0.0', '0.0', '0.0', '0.0']

$ python control_test_one_img.py -p AAA.pth -i ./Img/t_002_4.png
Model loaded from AAA.pth
予測された数字は: 4
各クラスの確率: ['0.0', '0.0', '0.0', '0.0', '1.0', '0.0', '0.0', '0.0', '0.0', '0.0']

$ python control_test_one_img.py -p AAA.pth -i ./Img/t_003_8.png
Model loaded from AAA.pth
予測された数字は: 8
各クラスの確率: ['0.0', '0.0', '0.0', '0.0', '0.0', '0.0', '0.0', '0.0', '1.0', '0.0']

$ python control_test_one_img.py -p AAA.pth -i ./Img/t_004_7.png
Model loaded from AAA.pth
予測された数字は: 7
各クラスの確率: ['0.0', '0.0', '0.0', '0.0', '0.0', '0.0', '0.0', '1.0', '0.0', '0.0']

$ python control_test_one_img.py -p AAA.pth -i ./Img/t_005_5.png
Model loaded from AAA.pth
予測された数字は: 5
各クラスの確率: ['0.0', '0.0', '0.0', '0.0', '0.0', '1.0', '0.0', '0.0', '0.0', '0.0']

かなり怪しい 6を描いたが、

$ python control_test_one_img.py -p AAA.pth -i ./Img/t_006_6.png
Model loaded from AAA.pth
予測された数字は: 6
各クラスの確率: ['0.0', '0.0', '0.0', '0.0', '0.0', '0.0', '0.8', '0.0', '0.2', '0.0']

---