推論の実行
次は、モデルを使って推論を行うコードを実装します。これには、最終的に推論コンテナ内で動作するコードを解説します。
1. モデルの読み込み
まず、学習済みモデルはPyTorchライブラリを使用して読み込みます。
model = torch.load(model_path + "/best.pt")
torch.load()はモデルを読み込むための関数です。学習結果の重みはbest.ptとして保存されているため、これを読み込みます。その際に、YOLOv5モデルの学習時にはmodelフォルダとutilフォルダも必要ですので、best.ptファイルと同じ階層に用意しておきます。この構造を保持しないと、モデルの読み込み時にエラーが発生しますので注意が必要です。
2. データの前処理
次に、Pillowというライブラリを使用して、画像をモデルに入力する前の前処理を行います。画像はLINE Botからbase64形式でエンコードされて送られてくるため、それをデコードします。
# nparrayに変換 image = base64.b64decode(flask.request.data) image = np.frombuffer(image, dtype="uint8") image = cv2.imdecode(image, cv2.IMREAD_COLOR) image = Image.fromarray(image)
YOLOv5モデルは640サイズの画像を入力とするため、このサイズに変更し、元画像の縦横比を保つために余白を追加します。
# リサイズ、余白追加
image_padding = Image.new('RGB', [max(image.size)] * 2, (0, 0, 0))
image_padding.paste(image, (0, 0, image.width, image.height))
image_resized = image_padding.resize((640, 640), Image.ANTIALIAS)
numpy配列に変換してからPyTorchのtensor型に変換します。
image = np.array(image_resized) image = torch.from_numpy(image)
数値を0～1の範囲に正規化します。
image = image.float() image /= 255
次の処理が特徴的で、モデルに入力するtensorのサイズを(1, 3, 640, 640)にリサイズします。通常、Pillowなどで読み込んだ際は(640, 640, 3)となっているので、これをモデルに入力できるよう整えていきます。
if len(image.shape) == 3:
image = image[None]
image = torch.permute(image, (0, 3, 1, 2))
これで前処理が完了しましたので、推論を実行してみましょう！
3. 推論の実行
PyTorchで読み込んだモデルはディクショナリー型になっているため、以下のように使います。
pred = model['model'](image)
4. モデルの出力
次に、推論結果がどのようになっているか確認しましょう。まず、得られた結果のサイズは以下のようになっています。
([[1, 25200, 6], [[1, 3, 80, 80, 6], [1, 3, 40, 40, 6], [1, 3, 20, 20, 6]]])
このデータの中身は次のようになります。
[ 3.99102e-01, 5.78214e-01, 1.90415e+00, -5.36558e-01, -8.57918e+00, 1.00173e+01], …
各6つの要素は次を表しています。
[center_x, center_y, width, height, 確度, 識別クラス係数]
数値の並びから、xとy座標が画像の左上から右下に向かって少しずつずれていくのが分かります。これは画像内の25,200個のバウンディングボックスに対して、推論の確度が示されていることを意味しています。また、今回のモデルはナンバーだけを検知するもので、識別クラスは1つしかありませんが、複数の物体を検知するモデルだとそれに応じて要素数が増加します。つまり、これらのデータを適切に処理して、推論結果を導き出していく必要があります。
5. 推論結果の導出
推論結果の導出にはNon-Maximum Suppression（NMS）という手法を使います。NMSの詳細な説明は省略しますが、複数のバウンディングボックスが重複している場合に、それを適切に処理して一つのバウンディングボックスとしてまとめる処理を行っています。このNMSの実装に関しては、YOLOv5のGitHubに掲載されているコードを参考にしました。
pred = non_max_suppression(pred)
print("処理後: ", pred)
処理後: [tensor([[292.75226, 326.84778, 355.07587, 358.43054, 0.93183, 0.00000],
[561.17792, 160.74695, 582.67120, 173.13055, 0.88415, 0.00000]])]
NMSを適用した結果はこのような形になっています。これらの要素は以下の情報を表します。
[top_left_x, top_left_y, bottom_right_x, bottom_right_y, 確度, クラス]
今回の例では、2つのナンバープレートが検知されました。これらを元に、検知したナンバープレートを黒で塗りつぶす処理を行います。
# ナンバー黒塗り for det in pred[0]: det = det * scale_ratio image_original = np.array(image_original) cv2.rectangle(image_original, pt1=(int(det[0]), int(det[1])), pt2=(int(det[2]), int(det[3])), color=(0, 0, 0), thickness=-1)
モデルへの入力画像は640x640でしたが、元画像のサイズに合わせてスケーリングを行っています。これでPythonコード上から推論を行う準備が整いました。次に、このモデルと処理をAWS環境にデプロイしていくことになります。
