PyTorchで学習したモデルをOpenCVで使ってみる

前回はPyTorchでマスクを付けている人を分類するモデルを作りました。

学習させたモデルをOpen CVで使ってみたいと思うのですが、毎回PyTorchを使うのは面倒なのでOpen CVだけでモデルを使う方法を試してみました。

モデルの変換

学習したモデルはPTHで保存できますが、これだと毎回PyTorchを呼び出す必要があります。

そこでモデルをOpen CVで使うためにOpen CVのDNNモジュールに対応しているONNX形式のモデルに変換する必要があります。

PyTorchではtorch.onnx.exportを使えば学習したモデルを簡単にONNX形式で保存できます。

また、Open CVのDNNモジュールを使う場合は、CPUを使うことが多いのでGPUで学習していた場合は念の為モデルをCPU向けに変換します。

これを踏まえて前回のモデルをONNXで変換する場合は以下の通りになります。

torch.onnx.export(model.to('cpu'), torch.randn(1, 3, 224, 224) , "face_mask_detector.onnx")

これで汎用性のあるモデルができました。

あとはDNNモジュールを使って以下のようにモデルを呼び出せば使えます。

faceMaskNet = cv2.dnn.readNetFromONNX("face_mask_detector.onnx")