[Amazon SageMaker] イメージ分類のモデルをNeoで最適化して、RasPi4+OpenCV+Webカメラで使用してみました

1 はじめに

CX事業本部の平内（SIN）です。

Amazon SageMager Neo（以下、Neo）を使用すると、既存のモデルを簡単にRaspberryPiで動作させる事ができます。
[Amazon SageMaker] イメージ分類のモデルをNeoで最適化してRaspberryPi Model 4で使用してみました

今回、DLRを利用する部分は、上記と同じですが、カメラをWebカメラに変更して、OpenCVの画像を入力として扱う要領を試してみました。

はじめに動作している様子です。

2 入力データ

Amazon SageMakerのイメージ分類（組み込みアルゴリズム）の入力となる、{"data":[1,3,224,224]} の形式にOpenCVの画像データを使用するためには、以下の点に考慮が必要です。

(1) 正方形

カメラから取得できる画像のサイズは、カメラの出力に依存しますが、通常、これは長方形になっていると思います。 このため、まず、正方形に変形（トリミング）する必要があります。

下記のコードでは、フレームが横長である場合、高さを１辺の長さとした正方形にトリミングしています。

img = img[0 : int(height), 0 : int(height)] 

(2) サイズ

サイズを244*244に変更する場合、resize() が利用可能です。正方形以外の画像で、この変換を行うと、画像が変形してしまうので注意が必要です。

img = cv2.resize(img, dsize=(224, 224)) 

(3) RGB

色相チャンネルのRBGの順は、期待されるもの違います。このため、下記のように変換が必要です。

# BGR => RGB
img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR) 

(4) ndarray

モデルの入力は、numpy.ndarrayです。OpenCVの画像データは、元々、numpy.ndarrayになっているので、この意味では変換の必要はありませんが、期待される、shapeが、(チャンネル数、列、行)であるため、この順番を入れ替える必要がります。

# (244,244,3) => (3,244,244)
img = img.transpose((2, 0, 1)) 

3 コード

使用したすべてのコードは以下のとおりです。推論に時間を要するため、2秒に１回だけ推論して、結果表示を更新しています。具体的には、カメラのフレームレートを5fpsに設定し、10フレームに１回だけ推論処理しています。

import numpy as np
from dlr import DLRModel
from PIL import Image
import cv2

# Webカメラ
DEVICE_ID = 0 

WIDTH = 800
HEIGHT = 600
FPS = 5

# Model
CLASSES = ['PORIPPY(GREEN)', 'OREO', 'CUNTRY_MAM', 'PORIPPY(RED)', 'BANANA', 
           'CHEDDER_CHEESE', 'PRETZEL(YELLOW)', 'FURUGURA(BROWN)', 'NOIR', 'PRIME',
           'CRATZ(RED)', 'CRATZ(GREEN)', 'PRETZEL(BLACK)', 'CRATZ(ORANGE)', 'ASPARA',
           'FURUGURA(RED)', 'PRETZEL(GREEN)']

model_path = './models'

class Model:
    def __init__(self, model_path):
        self.__model = DLRModel(model_path, 'cpu')

    def inference(self, image):
        out = self.__model.run({'data': image})
        return out

def putText(frame, x, y, text):
    font_size = 2
    font_color = (255, 255, 255)
    cv2.putText(frame, text, (x, y), cv2.FONT_HERSHEY_SIMPLEX, font_size, font_color, 3, cv2.LINE_AA)
    return frame

def decode_fourcc(v):
        v = int(v)
        return "".join([chr((v >> 8 * i) & 0xFF) for i in range(4)])

def main():

    # Model初期化
    model = Model(model_path)

    # カメラ初期化
    cap = cv2.VideoCapture (DEVICE_ID)

    # フォーマット・解像度・FPSの設定
    cap.set(cv2.CAP_PROP_FOURCC, cv2.VideoWriter_fourcc('Y','U','Y','V'))
    cap.set(cv2.CAP_PROP_FRAME_WIDTH, WIDTH)
    cap.set(cv2.CAP_PROP_FRAME_HEIGHT, HEIGHT)
    cap.set(cv2.CAP_PROP_FPS, FPS)

    # フォーマット・解像度・FPSの取得
    fourcc = decode_fourcc(cap.get(cv2.CAP_PROP_FOURCC))
    width = cap.get(cv2.CAP_PROP_FRAME_WIDTH)
    height = cap.get(cv2.CAP_PROP_FRAME_HEIGHT)
    fps = cap.get(cv2.CAP_PROP_FPS)
    print("fourcc:{} fps:{}　width:{}　height:{}".format(fourcc, fps, width, height))

    # 推論は2秒に１回実行する
    counter = 0
    interval = 2 

    name = ""
    probability = ""

    while True:
        # カメラ画像取得
        _, frame = cap.read()
        if(frame is None):
            continue

        frame = frame[0 : int(height), 0 : int(height)] # => 正方形

        counter += 1
        if(counter % (fps * interval) == 0):
            # ModelのInput用に画像変換
            image = cv2.resize(frame, dsize=(224, 224)) # => 244*244
            image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR) # RGB => BGR
            image = image.transpose((2, 0, 1)) # transpose shape (244,244,3)=>(3,244,244)

            # 推論
            out = model.inference(image)
            print(out)

            name = CLASSES[np.argmax(out[0])]
            probability = "{:.5f}".format(np.max(out))
            print("name:{} probability:{}".format(name,probability))

        # 結果（テキスト）表示
        frame = putText(frame, 10, int(height) - 100, name)
        frame = putText(frame, 10, int(height) - 30, probability)

        # 画像表示
        cv2.imshow('frame', frame)

        if cv2.waitKey(1) & 0xFF == ord('q'):
            break

    # VideoCaptureオブジェクト破棄
    cap.release()
    cv2.destroyAllWindows()

if __name__ == '__main__':
    main()

4 最後に

今回は、Neoで生成したモデルをOpenCVで取得した画像で試してみました。

SageMakerのエンドポイントと違い、入力データの形式を自前で整合させる必要があるので、少し、手間がかかると思います。