【現場で動くAI ④】重い環境と決別！ONNX変換でラズパイを高速化

ステップ4：エッジデバイスへの実装と最適化（ONNX変換）

前のステップ（ステップ3）でできるようになったこと：
PyTorchで転移学習を使い、MobileNet V3ベースのOK/NG分類モデルを構築・学習できるようになりました。ただし、このモデルを動かすにはPyTorch環境（数GB）が必要です。

■ 目的とゴール

ステップ3で学習したPyTorchのAIモデルを、フレームワークに依存しない世界共通のフォーマットである ONNX（オニキス） に変換します。これにより、現場のRaspberry Pi 5などの小型端末に重たいPyTorch環境をインストールすることなく、超軽量な実行エンジン（ONNX Runtime）だけで高速にOK/NG判定を行える仕組みを構築します。

今回は、諸事情によりRaspberry Piでテストしていますが、Macに特化したものを使っていないのでRaspberry Piでも問題なく動作する想定です。

■ 環境構築

ステップ0で作成した仮想環境を使います。

cd factory_ai_tutorial
source venv/bin/activate

# ONNX変換・実行ライブラリを追加インストール
pip install onnx onnxruntime onnxscript matplotlib

■ なぜONNXに変換するのか

PyTorchをそのままエッジデバイスで使わない理由：
PyTorchはライブラリのサイズが数GBに及ぶため、工場のラインに置くような安価で小さな端末の限られたメモリやストレージを圧迫します。またバージョンの依存関係によるトラブルも起きやすくなります。

ONNXのメリット：
学習済みモデルをONNX形式（.onnx）に書き出すことで、実行環境からPyTorchを完全に切り離せます。エッジ側では、ステップ1のOpenCVと、推論専用の軽量エンジン ONNX Runtime さえあればAIを動かすことが可能です。

■ サンプルプログラム

1. モデルのONNX変換スクリプト（step4_export.py）

import torch
import torch.nn as nn
from torchvision import models

# 1. 変換元となるPyTorchモデルの準備（ステップ3と同じ構成）
model = models.mobilenet_v3_small(weights=models.MobileNet_V3_Small_Weights.DEFAULT)
num_features = model.classifier[3].in_features
model.classifier[3] = nn.Linear(num_features, 2)

# ※実運用時はここで学習済みの重み（.pth）をロードします
# model.load_state_dict(torch.load('factory_model.pth'))

model.eval()

# 2. ONNX形式へのエクスポート
dummy_input = torch.randn(1, 3, 224, 224)

onnx_file_path = "factory_model.onnx"
torch.onnx.export(
    model,
    dummy_input,
    onnx_file_path,
    export_params=True,
    opset_version=18,
    input_names=['input'],
    output_names=['output']
)

print(f"ONNXモデルを '{onnx_file_path}' として出力しました。")

2. エッジデバイスでの推論スクリプト（step4_inference.py）

現場のRaspberry Pi 5等で動かすためのコードです。コード内に import torch が一切ないことに注目してください。

import cv2
import numpy as np
import onnxruntime as ort
import matplotlib.pyplot as plt

# 1. ONNX Runtimeの準備
onnx_session = ort.InferenceSession("factory_model.onnx")

# 2. 画像の読み込みと前処理
img_bgr = cv2.imread('sample.jpg')
if img_bgr is None:
    print("エラー: sample.jpgが見つかりません。")
    exit()

img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
img_resized = cv2.resize(img_rgb, (224, 224))

# 0〜255のピクセル値を0.0〜1.0に変換
img_normalized = img_resized.astype(np.float32) / 255.0

# --- ImageNet学習済みモデルが要求する標準化（Normalize）---
mean = np.array([0.485, 0.456, 0.406])
std  = np.array([0.229, 0.224, 0.225])
img_normalized = (img_normalized - mean) / std

# PyTorchの入力形式 (Batch, Channel, Height, Width) に並び替え
img_transposed = np.transpose(img_normalized, (2, 0, 1))
input_data = np.expand_dims(img_transposed, axis=0).astype(np.float32)

# 3. AIによる推論の実行
outputs = onnx_session.run(None, {'input': input_data})
predictions = outputs[0]  

probs = np.exp(predictions) / np.sum(np.exp(predictions), axis=1, keepdims=True)
class_id = np.argmax(probs)
confidence = np.max(probs)

labels = ["OK (Normal)", "NG (Defect)"]
print(f"判定結果: {labels[class_id]} (確信度: {confidence * 100:.1f}%)")

# ==========================================
# 4. 結果を画像に描画して表示 (★ここから下を書き換え)
# ==========================================
# 文字入れ（OpenCVで画像に直接書き込む）
cv2.putText(img_bgr, f"{labels[class_id]} ({confidence * 100:.1f}%)",
            (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1.0, (0, 255, 0), 2)

# OpenCVのBGR形式から、Matplotlib表示用のRGB形式に変換
img_result_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)

# Matplotlibで安全に表示
plt.figure(figsize=(6, 6))
plt.imshow(img_result_rgb)
plt.title(f"Result: {labels[class_id]} ({confidence * 100:.1f}%)")
plt.axis('off') # グラフの目盛りを消す
plt.show()      # ここでウィンドウが開き、バツボタンで安全に閉じられます

■ プログラムの実行と確認

変換スクリプトを実行してONNXファイルを生成

python step4_export.py

実行結果

[torch.onnx] Obtain model graph for `MobileNetV3([...]` with `torch.export.export(..., strict=False)`...
[torch.onnx] Obtain model graph for `MobileNetV3([...]` with `torch.export.export(..., strict=False)`... ✅
[torch.onnx] Run decompositions...
/Users/nomura/.pyenv/versions/3.13.2/lib/python3.13/copyreg.py:99: FutureWarning: `isinstance(treespec, LeafSpec)` is deprecated, use `isinstance(treespec, TreeSpec) and treespec.is_leaf()` instead.
  return cls.__new__(cls, *args)
[torch.onnx] Run decompositions... ✅
[torch.onnx] Translate the graph into ONNX...
[torch.onnx] Translate the graph into ONNX... ✅
[torch.onnx] Optimize the ONNX graph...
[torch.onnx] Optimize the ONNX graph... ✅
ONNXモデルを 'factory_model.onnx' として出力しました。

推論スクリプトを実行（sample.jpgが必要）

python step4_inference.py

実行結果

判定結果: NG (Defect) (確信度: 67.5%)

■ まとめと次のステップへ

この factory_model.onnx と step4_inference.py をRaspberry Pi 5にコピーすれば、そのまま工場の現場で動くAIシステムになります。エッジデバイスでのAI判定の最大のメリットは「ネットワークに繋がっていなくても自己完結して判定できる」点です。判定結果をGPIOピンから出力してパトライト（回転灯）や排除アームと連動させるなど、物理的なシステムとも簡単に連携できます。

次はこのシステムを実際の工場ラインに設置して検証する「ステップ5：PoCとビジネス価値の証明」へ進みます。