OpenCV4Python10:OpenCV(numpy.ndarray)とPyTorch Hubで画像分類
【0】はじめに
Pytorchにも学習済みモデルを置いたモデルリポジトリとして「PyTorch Hub」がある。
今回は「PyTorch HubのResNetモデル」+「OpenCV」で画像分類してみる
【1】PyTorch Hub:ResNetモデルのロード
PyTorch Hub:ResNetモデルは以下。
サイトの例に書いてあるように「torch.hub.load」を使ってモデルをロードすればよい。
【例1】:PyTorch HubからResNetモデルをロードする
import torch
... ...
# モデルのロード
model = torch.hub.load('pytorch/vision:v0.8.2', 'resnet18', pretrained=True)
model.eval() # 評価モードに切り替える
公式サイトの例ではバージョン「v0.6.0」という指定をしていたが、最新のバージョンはGithubから確認できる。(※v0.8.2があったのでこれを今回設定してみている)
あとはOpenCVで画像を読み込んで前処理実施して学習済みモデルにデータを渡して画像分類させる。
以下画像を読み込んでみる。(※サンプルではurllibで画像取得をしているが、ここではダウンロードしてファイルとして読み込む)
(https://github.com/pytorch/hub/raw/master/images/dog.jpg)
【2】前処理:OpenCVとtorchvision.transforms
ResNetモデルでは画像データを以下のように変換する「前処理」が必要。
・形状(3x224x224) → モデルに渡すときは[1, 3, 224, 224]にする
・画素値:0~1.0
・正規化:平均値= [0.485, 0.456, 0.406]、標準偏差= [0.229, 0.224, 0.225]にならす(学習データであるImageNet画像の画像データの分布に合わせる)
■OpenCVで画像の読み込み
OpenCVで普通に読み込んだ後チャネル順はRGBにしておく。
【例2】:画像の読み込み
import cv2 as cv
import numpy as np
... ...
input_image = cv.imread('dog.jpg',cv.IMREAD_COLOR)
input_image = cv.cvtColor(input_image,cv.COLOR_BGR2RGB)
■前処理:torchvision.transforms
OpenCVの画像データ(numpy.ndarray)の範囲で何とかしようとするよりもPyTorchでは前処理として「torchvision.transforms」使うのが楽。
具体的には「torchvision.transforms.Compose」でいくつかの前処理をまとめる。
【例3】:transforms.Composeで前処理を設定
from torchvision import transforms
... ...
# 前処理の定義 (公式サンプルの前処理に追記)
preprocess = transforms.Compose([
transforms.ToPILImage(), # numpy.ndarray→Pillowオブジェクトに変換
transforms.Resize(256), # 画像サイズを256に変換
transforms.CenterCrop(224), # 画像サイズ224でセンタークロッピング
transforms.ToTensor(), # torch.Tensorオブジェクト化
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), # 正規化
])
# OpenCVで読み込んだ画像データ(numpy.ndarray)を前処理を適用する
input_tensor = preprocess(input_image)
torchvision.transformsの各処理で受け付けるオブジェクトの型は「PIL(Pillow)オブジェクト」と「torch.Tensorオブジェクト」であることに注意する。
そのためOpenCVで読み込んだ画像データ(numpy.ndarray)はどこかで型変換する必要がある。今回は「transforms.Compose」の最初に「ToPILImage()」を配置して、「numpy.ndarray」→ 「PIL(Pillow)オブジェクト」に変換している。
【例4】画像データ形状の確認
前処理によって返ってきたtorch.Tensorオブジェクトの形状を確認。
print(input_tensor.shape)
出力結果:
torch.Size([3, 224, 224])
モデルにあたえる画像データの形状は[1,3,224,224]とする必要があるため、torch.Tensorオブジェクトの形状を変換する。
変換には「unsqueeze()」を使う。
【例5】:torch.Tensorの形状を(3, 224, 224)→(1, 3, 224, 224)にする
# torch.Tensorオブジェクトの形状を変換(次元を追加)
input_batch = input_tensor.unsqueeze(0)
print(input_batch.shape)
出力結果:
torch.Size([1, 3, 224, 224])
あとはモデルに渡して画像分類をすればよい。
【3】画像分類をする
PyTorchでモデルに新たなデータを評価させる時は、モデルのモードに気をつける。「eval()」で評価モードにしていること(今回はモデルロード時に設定済み)と、「no_grad()」で勾配計算を行わないようにする。
【例6】:OpenCVで読み込んだ画像をモデルに渡して分類する
... ...
model = torch.hub.load('pytorch/vision:v0.8.2', 'resnet18', pretrained=True)
model.eval() # 評価モードに切り替える
... ...
# GPU(cuda)が使える場合は利用する設定
if torch.cuda.is_available():
input_batch = input_batch.to('cuda')
model.to('cuda')
# model.eval()とセットでno_gradで勾配計算はOFFにして予測させる
with torch.no_grad():
output = model(input_batch)
print(output[0]) # output[0]→ 1枚目の画像の分類結果
出力結果:
tensor([ 1.5917e-02, -1.5497e+00, 3.2031e-01, -2.0585e+00, -8.5747e-01,
1.7843e+00, 1.4699e+00, 2.1626e+00, 4.4888e+00, 8.2885e-01
... ...
... ...
-6.5606e-01, -1.8088e+00, -2.9126e+00, 5.6032e-01, 2.5117e+00],
device='cuda:0')
あとは出力値の調整をして、あらかじめ用意しておいたimagenetの分類項目を記載したテキストファイル
https://raw.githubusercontent.com/pytorch/hub/master/imagenet_classes.txt
と照らし合わせればよい。
公式のサンプルのままだが、softmax関数でスケーリングしてトップ5を取得、テキストファイルと突き合わせる。
【例7】:モデルの結果とimagenetの項目テキストを突き合わせる
probabilities = torch.nn.functional.softmax(output[0], dim=0)
probabilities
# カテゴリファイルを読み込む
with open("imagenet_classes.txt", "r") as f:
categories = [s.strip() for s in f.readlines()]
# トップ5の値とインデックスを取得
top5_prob, top5_catid = torch.topk(probabilities, 5)
# カテゴリファイルと突き合わせて分類項目を確認する
for i in range(top5_prob.size(0)):
print(categories[top5_catid[i]], top5_prob[i].item())
出力結果:
Samoyed 0.8846219182014465
Arctic fox 0.0458051860332489
white wolf 0.04427620768547058
Pomeranian 0.005621347576379776
Great Pyrenees 0.00465201074257493
▲ResNetでは「サモエド(犬)」である可能性がかなり高いと判断
【全体コード】
「dog.jpg」と「imagenet_classes.txt」は手元にダウンロードしている前提。
import torch
from torchvision import transforms
import cv2 as cv
import numpy as np
# モデルのロード
model = torch.hub.load('pytorch/vision:v0.8.2', 'resnet18', pretrained=True)
model.eval() # 評価モードに切り替える
# OpenCVで画像の読み込み
filename = 'dog.jpg'
input_image = cv.imread(filename,cv.IMREAD_COLOR)
input_image = cv.cvtColor(input_image,cv.COLOR_BGR2RGB) # チャネル順の変換BGR⇒RGB
# 前処理定義
preprocess = transforms.Compose([
transforms.ToPILImage(), # numpy.ndarray→Pillowオブジェクトに変換
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
# 読み込んだ画像の前処理実行
input_tensor = preprocess(input_image)
#print(input_tensor.shape) # [output]:torch.Size([3, 224, 224])
# torch.Tensorの形状変換
input_batch = input_tensor.unsqueeze(0)
#print(input_batch.shape) # [output]:torch.Size([1, 3, 224, 224])
# GPU(cuda)が使える場合は利用する設定
if torch.cuda.is_available():
input_batch = input_batch.to('cuda')
model.to('cuda')
# model.eval()とセットでno_gradで勾配計算はOFFにして予測させる
with torch.no_grad():
output = model(input_batch)
#print(output[0]) # output[0]→ 1枚目の画像の分類結果
# softmax関数で出力結果をスケーリング
probabilities = torch.nn.functional.softmax(output[0], dim=0)
print(probabilities)
# imagenetの項目テキストと突き合わせ
# カテゴリファイルを読み込む
with open("imagenet_classes.txt", "r") as f:
categories = [s.strip() for s in f.readlines()]
# トップ5の値とインデックスを取得
top5_prob, top5_catid = torch.topk(probabilities, 5)
# カテゴリファイルと突き合わせて分類項目を確認する
for i in range(top5_prob.size(0)):
print(categories[top5_catid[i]], top5_prob[i].item())
【おまけ】
ほとんどPyTorchの説明になってしまったが、まとめると今回のOpenCVの出番は画像の読み込みだけ。前処理はPytorchの「torchvision.transforms」に渡してしまうためほぼ出番はない。
さらにいうとPyTorchにも画像を読み込むライブラリは用意されている。
OpenCVを使わないとできないような処理を組む予定がない、画像読み込みだけ、というのであれば、このライブラリでも十分。