[OpenCV] Pytorchの手書き数字（MNIST）分類モデルをOpenCVから利用してみました

1 はじめに

CX事業本部の平内（SIN）です。

Pytorch入門ということで、MNIST（手書き数字のデータセット）から作成したモデルを使用して、OpenCVでWebカメラの動画を推論にかけてみました。

使用したモデルのコードは、Githubで公開されている、Pytorchの公式サンプルコードです。
https://github.com/pytorch/examples/blob/master/mnist/main.py

最初に動作しているようすです。 推論の対象となっているのは、画面の中央だけで、別ウインドウに表示されている部分です。

2 モデル

公開されているサンプルコードは、以下のようになっています。

(1) データセット

データセットは、torchvisionによって、MNISTが利用されています。 取得時に、transformsによる変換を行って、訓練用とテスト用のデータローダーが準備されています。

ちなみに、コマンドラインからデフォルト値で使用した場合、学習用のミニバッチは、64となります。

from torchvision import datasets, transforms
transform=transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,))
        ])
dataset1 = datasets.MNIST('../data', train=True, download=True, transform=transform)
dataset2 = datasets.MNIST('../data', train=False, transform=transform)
train_loader = torch.utils.data.DataLoader(dataset1, **kwargs)
test_loader = torch.utils.data.DataLoader(dataset2, **kwargs)

(2) モデル

入力は、畳み込み層✕2、プーリング層✕1で処理された後、全結合層（9216 -> 128 -> 10）を通過します。また、途中で、ドロップアウトも挟まれています。

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, 3, 1)
        self.conv2 = nn.Conv2d(32, 64, 3, 1)
        self.dropout1 = nn.Dropout2d(0.25)
        self.dropout2 = nn.Dropout2d(0.5)
        self.fc1 = nn.Linear(9216, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = self.conv1(x)
        x = F.relu(x)
        x = self.conv2(x)
        x = F.relu(x)
        x = F.max_pool2d(x, 2)
        x = self.dropout1(x)
        x = torch.flatten(x, 1)
        x = self.fc1(x)
        x = F.relu(x)
        x = self.dropout2(x)
        x = self.fc2(x)
        output = F.log_softmax(x, dim=1)
        return output

(3) 学習

オプティマイザは、optim.Adadelta()で、学習レート(lr)は、デフォルト値で1.0となっています。 また、損失関数 F.nll_loss()です。

optimizer = optim.Adadelta(model.parameters(), lr=args.lr)
F.nll_loss(output, target)

下記は、１エポックの中で呼ばれている訓練用のコードです。 ミニバッチごとにデータローダから訓練用データを取り出し、学習が進められています。

def train(args, model, device, train_loader, optimizer, epoch):
    model.train()
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        output = model(data)
        loss = F.nll_loss(output, target)
        loss.backward()
        optimizer.step()
        if batch_idx % args.log_interval == 0:
            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                epoch, batch_idx * len(data), len(train_loader.dataset),
                100. * batch_idx / len(train_loader), loss.item()))
            if args.dry_run:
                break

コマンドラインから、下記のように利用することで、エポック14回（デフォルト値）、学習が進みます。 手元のMBP（2.3GHz Core i5）では、1エポックに、約2分強で、14エポック終了まで約30分程度かかりました。

テストデータの検証結果を見た感じ、epochは、7回ぐらいで良いのかも知れません。

% python3 index.py --log-interval=100 --save-model
Train Epoch: 1 [0/60000 (0%)]   Loss: 2.293032
Train Epoch: 1 [6400/60000 (11%)]   Loss: 0.273636
Train Epoch: 1 [12800/60000 (21%)]  Loss: 0.294438
Train Epoch: 1 [19200/60000 (32%)]  Loss: 0.360903
Train Epoch: 1 [25600/60000 (43%)]  Loss: 0.097714
Train Epoch: 1 [32000/60000 (53%)]  Loss: 0.270267
Train Epoch: 1 [38400/60000 (64%)]  Loss: 0.145856
Train Epoch: 1 [44800/60000 (75%)]  Loss: 0.173067
Train Epoch: 1 [51200/60000 (85%)]  Loss: 0.289421
Train Epoch: 1 [57600/60000 (96%)]  Loss: 0.154544

Test set: Average loss: 0.0497, Accuracy: 9835/10000 (98%)
・・・略・・・
Train Epoch: 2
Test set: Average loss: 0.0409, Accuracy: 9871/10000 (99%)
Train Epoch: 3
Test set: Average loss: 0.0354, Accuracy: 9876/10000 (99%)
Train Epoch: 4
Test set: Average loss: 0.0342, Accuracy: 9886/10000 (99%)
Train Epoch: 5
Test set: Average loss: 0.0325, Accuracy: 9893/10000 (99%)
Train Epoch: 6
Test set: Average loss: 0.0317, Accuracy: 9896/10000 (99%)
Train Epoch: 7
Test set: Average loss: 0.0303, Accuracy: 9900/10000 (99%)
Train Epoch: 8
Test set: Average loss: 0.0308, Accuracy: 9901/10000 (99%)
Train Epoch: 9
Test set: Average loss: 0.0298, Accuracy: 9905/10000 (99%)
Train Epoch: 10
Test set: Average loss: 0.0294, Accuracy: 9910/10000 (99%)
Train Epoch: 11
Test set: Average loss: 0.0292, Accuracy: 9910/10000 (99%)
Train Epoch: 12
Test set: Average loss: 0.0287, Accuracy: 9911/10000 (99%)
Train Epoch: 13
Test set: Average loss: 0.0287, Accuracy: 9910/10000 (99%)
Train Epoch: 14
Test set: Average loss: 0.0286, Accuracy: 9908/10000 (99%)

--save-modelのスイッチを付けておくと、学習終了時にモデル（mnist_cnn.pt）が保存されます。

 % ls -la *.pt
-rw-r--r--  1 hirauchi.shinichi  staff  4800893  7 19 02:58 mnist_cnn.pt

3 入力形式

モデルの入力形式を確認するために、データローダからデータを1つ取り出して確認しました。

import torch
from torchvision import datasets, transforms
import numpy as np
import matplotlib.pyplot as plt

transform=transforms.Compose([  transforms.ToTensor() ])
test_data = datasets.MNIST('../data', train=False, download=True, transform=transform)
test_loader = torch.utils.data.DataLoader(test_data, batch_size=1, shuffle=False)

tmp = test_loader.__iter__()
data, labels = tmp.next() 

# size
print(data.size()) # torch.Size([1, 1, 28, 28])

# imshow
img = data.numpy().reshape((28, 28))
plt.imshow(img, cmap='gray')

# dump
print(data)

(1) size()

学習時にモデルに送られるデータは、4次元のTorch型でtorch.Size([1, 1, 28, 28])となっています。

(2) imshow()

numpyに変換してmatplotlibで表示すると以下のように、黒字に白で数字が書かれたデータであることが分かります。

(3) dump

ダンプしてみると、黒が0で白が1となる範囲で、表現されていることを確認できます。

tensor([[[[0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
           0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
           0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
           0.0000, 0.0000, 0.0000, 0.0000],
・・・省略・・・
          [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.3294, 0.7255,
           0.6235, 0.5922, 0.2353, 0.1412, 0.0000, 0.0000, 0.0000, 0.0000,
           0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
           0.0000, 0.0000, 0.0000, 0.0000],
・・・省略・・・
          [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
           0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
           0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
           0.0000, 0.0000, 0.0000, 0.0000]]]])

OpenCVで取得した画像を、上記の入力形式に合わせてやることで、モデルが利用可能になります。

4 使ってみる

モデルを使っているコードです。

OpenCVで取得した画像は、下記の変換を行って、モデルの入力としています。

• 推論の範囲を切り出す
• グレーススケールで白黒画像に変換する
• 白と黒を反転させる
• 2値化処理を行う
• サイズを28✕28に変換する
• 0〜255で表現されているデータを0〜1に変換する
• 次元を追加して、[1, 1, 28,28]に合わせる

import cv2
import time
import torch
import numpy as np
import torch.nn as nn
import torch.nn.functional as F

# Webカメラ
DEVICE_ID = 0 
WIDTH = 800
HEIGHT = 600

path = './mnist_cnn.pt'

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, 3, 1)
        self.conv2 = nn.Conv2d(32, 64, 3, 1)
        self.dropout1 = nn.Dropout2d(0.25)
        self.dropout2 = nn.Dropout2d(0.5)
        self.fc1 = nn.Linear(9216, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = self.conv1(x)
        x = F.relu(x)
        x = self.conv2(x)
        x = F.relu(x)
        x = F.max_pool2d(x, 2)
        x = self.dropout1(x)
        x = torch.flatten(x, 1)
        x = self.fc1(x)
        x = F.relu(x)
        x = self.dropout2(x)
        x = self.fc2(x)
        output = F.log_softmax(x, dim=1)
        return output

def main():
    cap = cv2.VideoCapture (DEVICE_ID)

    # フォーマット・解像度・FPSの設定
    cap.set(cv2.CAP_PROP_FRAME_WIDTH, WIDTH)
    cap.set(cv2.CAP_PROP_FRAME_HEIGHT, HEIGHT)

    # フォーマット・解像度・FPSの取得
    width = cap.get(cv2.CAP_PROP_FRAME_WIDTH)
    height = cap.get(cv2.CAP_PROP_FRAME_HEIGHT)
    fps = cap.get(cv2.CAP_PROP_FPS)
    print("fps:{}　width:{}　height:{}".format(fps, width, height))

    # 検出範囲
    size = 50
    x1 = int(width/2-size)
    x2 = int(width/2+size)
    y1 = int(height/2-size)
    y2 = int(height/2+size)

    # モデル＋パラメータの読み込み
    model = Net()
    model.load_state_dict(torch.load(path))
    model.eval() # 評価モード

    while True:

        # カメラ画像取得
        _, frame = cap.read()
        if(frame is None):
            continue

        img = frame[y1 : y2, x1 : x2] # 対象領域の抽出
        img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) # グレースケール
        img = cv2.bitwise_not(img) # 白黒反転
        _, img = cv2.threshold(img, 0, 255, cv2.THRESH_OTSU) # 2値化
        img = cv2.resize(img,(28, 28)) # サイズ変更 => 28 * 28
        cv2.imshow('img', img) # モニター

        img = img/256 # 0〜255 => 0.0〜1.0
        img = img[np.newaxis, np.newaxis, :, :] # 次元追加　(28,28) => (1, 1, 28, 28) 
        pred = model(torch.tensor(img, dtype=torch.float32))

        print(pred)
        index = pred.argmax()
        prediction = "Prediction: {}".format(pred.argmax().item())
        print(prediction)

        cv2.rectangle(frame, (x1, y1), (x2, y2), (255, 255, 0))
        cv2.putText(frame, prediction, (10, int(height)-50), cv2.FONT_HERSHEY_SIMPLEX, 3, (255, 255, 255), 3, cv2.LINE_AA)

        # 画像表示
        cv2.imshow('frame', frame)

        # 'q'をタイプされたらループから抜ける
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break

    # VideoCaptureオブジェクト破棄
    cap.release()
    cv2.destroyAllWindows()

if __name__ == '__main__':
    main()

5 最後に

学習されたモデルを使用するために、入力形式を整合させることは必須の作業です。

作業としては、概ね、「縦、横、チャンネル」の順番や、ミニバッチの分の次元調整が主になると思います。 また、OpenCVでは、Numpy形式で画像を扱いますが、最終的にはTorch型にすることも忘れてはなりません。

次元やデータサイズが、モデルの入力に一致しないと、推論の段階でエラーとなるはずなので、そのエラーメッセージを見ることでも、だいだい間違いに気がつくかも知れません。