OpenCV4Python10:OpenCV(numpy.ndarray)とPyTorch Hubで画像分類

【０】はじめに

Pytorchにも学習済みモデルを置いたモデルリポジトリとして「PyTorch Hub」がある。

PyTorch

今回は「PyTorch HubのResNetモデル」＋「OpenCV」で画像分類してみる

【１】PyTorch Hub：ResNetモデルのロード

PyTorch Hub：ResNetモデルは以下。

PyTorch

サイトの例に書いてあるように「torch.hub.load」を使ってモデルをロードすればよい。

torch.hub — PyTorch master documentation

【例１】：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)

【２】前処理：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にしておく。

【例２】：画像の読み込み

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 — PyTorch 1.7.1 documentation

具体的には「torchvision.transforms.Compose」でいくつかの前処理をまとめる。

torchvision.transforms — PyTorch 1.7.1 documentation

【例３】：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)オブジェクト」に変換している。

【例４】画像データ形状の確認
前処理によって返ってきたtorch.Tensorオブジェクトの形状を確認。

print(input_tensor.shape)

出力結果：

torch.Size([3, 224, 224])

モデルにあたえる画像データの形状は[1,3,224,224]とする必要があるため、torch.Tensorオブジェクトの形状を変換する。
変換には「unsqueeze()」を使う。

torch.unsqueeze — PyTorch 1.7.1 documentation

【例５】：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])

あとはモデルに渡して画像分類をすればよい。

【３】画像分類をする

PyTorchでモデルに新たなデータを評価させる時は、モデルのモードに気をつける。「eval()」で評価モードにしていること(今回はモデルロード時に設定済み)と、「no_grad()」で勾配計算を行わないようにする。

【例６】：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関数でスケーリングしてトップ５を取得、テキストファイルと突き合わせる。

torch.nn.functional — PyTorch 1.7.1 documentation

torch.topk — PyTorch 1.7.1 documentation

【例７】：モデルの結果と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()]

# トップ５の値とインデックスを取得
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()]

# トップ５の値とインデックスを取得
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にも画像を読み込むライブラリは用意されている。

torchvision.io — PyTorch 1.7.1 documentation

OpenCVを使わないとできないような処理を組む予定がない、画像読み込みだけ、というのであれば、このライブラリでも十分。