537 views

【1】やりたいこと

過去記事 (36)【PyTorchでMNIST #7】自筆画像PNGファイルを自動認識させる。
で作ったプログラムを、
過去記事 (39) 共用サーバでPyTorach（順伝播だけ）を実行したい。
の調査結果を使い Webサービス化したい。

つまり・・・
Webブラウザ上で数字を手書きし、
　その自筆数字を共用サーバに送信し、
　　学習済AI（PyTorch使用プログラム）で自動認識させたい。

【2】やってみる

1) プログラム構成と処理フロー

後日追記

2) プログラム・ソースコード

(1) dataset_MNIST.py

前回からの変更箇所をハイライト表示する。

import torch
from torchvision import datasets, transforms
from PIL import Image
import io
import base64

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] に拡張

    def load_one_image_from_base64(self, image_base64):
        # Web経由で送られてきたBase64エンコード画像データを、
        # MNIST準拠のTensor [1,1,28,28] に変換して返す。
        # Base64データのヘッダー部分（data:image/png;base64,）を取り除く
        if ',' in image_base64:
            image_base64 = image_base64.split(',')[1]
        # Base64をデコードしてバイナリに
        image_bytes = base64.b64decode(image_base64)
        # バイナリデータから画像を開く
        image = Image.open(io.BytesIO(image_bytes)).convert('L')  # モノクロ変換
        # サイズを28x28に統一
        image = image.resize((28, 28))
        # Tensor化＋正規化
        image = self.transform(image)
        # 形状を [1,1,28,28] に変換して返す
        return image.unsqueeze(0)

(2) control_test_one_img.py

既存の predicti_image を分解し、
・predict_image_from_file : 画像ファイル入力用
・predict_image_from_base64 : Web受信データ入力用
に分けて実装した。
前回からの変更箇所をハイライト表示する。

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

#//////////////////////////////////////////////////////////////////////////
def predict_image_from_file(img_path, model_path):
    # 画像読み込みと前処理
    dataset = MyDataset()
    image_tensor = dataset.load_one_image(img_path)
    # 自動認識を実行
    return exec_predicti_image(image_tensor, model_path)

#//////////////////////////////////////////////////////////////////////////
def predict_image_from_base64(img_data, model_path):
    # 画像読み込みと前処理
    dataset = MyDataset()
    image_tensor = dataset.load_one_image_from_base64(img_data)
    # 自動認識を実行
    return exec_predicti_image(image_tensor, model_path)

#//////////////////////////////////////////////////////////////////////////
def exec_predicti_image(image_tensor, model_path):
    ret = {}
    # 環境設定
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    image_tensor = image_tensor.to(device)
    # モデルとパラメータの読み込み
    model = Net()
    trainer = Trainer(model, device)
    trainer.load(model_path)
    # 推論
    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]
        ret['label'] = predicted_label
        ret['probs'] = prob_array
    return ret

#//////////////////////////////////////////////////////////////////////////
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()

3) 実行結果

こちらに設置した。
この子は 6 の認識が苦手かもしれない・・・
https://www.dogrow.net/nnet/sample/00040/