はじめに

本記事では、Raspberry Pi AIカメラでオリジナルのキーポイント検出AIモデルを作成し、実際に動かすまでのステップを解説します。

Raspberry Pi AIカメラ上で動作可能なAIモデルを簡単に構築できるように、サンプルコードを公開しており、本記事はそのツールを使ってみて実際にAIモデルを作る流れを紹介します。

前回の記事①はNanodetを用いた物体検知モデルを紹介しました。

今回は第2弾としてキーポイント検出モデルを学習していきます。

🎯 この記事の概要

サンプルコードを使って、キーポイント検出モデル を作成する手順
実際に生成されるモデルファイルとその使い方の紹介

💻 キーポイントAIモデルとは

キーポイントAIモデルとは、人や物体の特徴的な点（キーポイント）を検出し、位置情報として出力するAIモデルです。

Raspberry Pi AIカメラ向けには、人体の骨格（関節）に特化した事前学習済みモデルが提供されており、姿勢推定や動作認識にそのまま利用できます。

🔹 公開済みの事前学習済みモデルでできること

人の全身の骨格点（頭・肩・肘・手首・膝・足首など）をリアルタイムに推定

🔹 公開モデルの利用方法

Raspberry Pi AIカメラ向けに提供される事前学習済みモデルは、以下の公式チュートリアルから取得できます。
これらは人体骨格に対して学習済みのため、追加学習なしで利用可能です。
公式チュートリアルはこちらになります。
https://www.raspberrypi.com/documentation/accessories/ai-camera.html

⭐️⭐️ 再学習による応用範囲の拡大(←本記事の紹介内容)

キーポイントAIモデルは、対象物とその特徴点の定義を変えることで、人以外の物体にも応用可能です。
たとえば再学習することで、以下のようなユースケースが実現できます。

ロボットの関節位置推定
動物の姿勢認識（犬や猫の動作解析など）
製造ラインのアナログゲージの読み取り

つまり、キーポイントモデルは「どの点を検出するか」を自由に定義できるため、
再学習次第で幅広いエッジAIアプリケーションに活用可能です。

本記事では、簡単に試せることができる一例として、矢印の図形の辺を検出するモデルを作成します。

もちろん、本スクリプトを使用してデータセットを置き換えることで任意の学習が可能になります。
今回はすぐ試すためにサンプルとして公開しているクリーンデータセットの矢印データを使用してみます。

🔧 環境要件

Raspberry Pi + AIカメラ
学習用に使用するPC
　- NVIDIA GPU推奨 (学習時)
　- OS: Ubuntu 22.04
　- Python 3.10

ローカルでの環境構築

ここではローカル環境に学習環境を構築する場合を紹介します。
すでに物体検知モデル編①にて、環境構築を実施済みの方はこちらのセクションはスキップしてください。

①　まずはリポジトリをクローンします。

git clone https://github.com/SonySemiconductorSolutions/aitrios-rpi-training-samples.git
cd aitrios-rpi-training-samples

②セットアップ
必要なパッケージをインストールします。

sudo apt update
sudo apt -y install --no-install-recommends
apt-utils build-essential libgl1 libgl1-mesa-glx
libglib2.0-0 python3 python3-pip python3-setuptools git gnupg

③Python3.10の仮想環境の構築・有効化
本チュートリアルではPython3.10を想定しているのでまずは3.10が入っていることを確認しましょう。

インストールするライブラリのバージョン依存関係があるので、必ずPython3.10系を使うようにしましょう

python3.10 --version

このように表示されれば、Python3.10がインストールされています。

まだの方はたとえば以下の手順でpython3.10をインストールします。
(すでにpython3.10をインストール済みの方はスキップしてください)

sudo apt update
sudo apt install -y software-properties-common
sudo add-apt-repository ppa:deadsnakes/ppa
sudo apt update
sudo apt install -y python3.10 python3.10-venv python3.10-dev

次に仮想環境を構築します。

python3.10 -m venv .venv
source .venv/bin/activate

④パッケージのインストール

pip install .
pip install -e third_party/nanodet/nanodet

Step 1: モデルを学習する

本リポジトリでは、AIモデルの学習・量子化・評価に関する各種設定を .ini ファイルで管理しています。
今回はこちらのposenet_arrow.iniのモデルを学習していきます。

✅ ポイント
各種設定を.ini ファイルで管理することで、なるべく設定値の変更だけで柔軟に学習条件を調整できるにしています。
iniファイルでは、使用するデータセットやタスクの選択から、パラメータの調整も可能です。
調整方法は記事後半の応用：自作データセットに置き換えるには？の章にて説明します。

使用するデータセット

今回の学習で使用するのは、矢印の図形のデータセットです。

ちなみにもともとの図形の画像はこちらで公開しています。

サンプルデータセットでは、それぞれの頂点に対して以下のように7つのポイントをつけています。

ではこれから、指定された .ini ファイルをもとに、モデルの学習および量子化を行います。

iniファイルが格納されているsamplesフォルダに移動し、コマンドを実行します。
以下のコマンドではAIモデルの学習からモデルの量子化まで自動で実行します。

cd samples

imx500_zoo posenet_arrow.ini

実行するとデータセットをDLしたのち、トレーニングが開始されていきます。
デフォルトの設定ファイルのままだと50epochが指定されています。

完了すると以下のようなモデルファイルが生成されます。
以下のファイル群が作成されていれば成功です。

./samples/model/posenet_arrow/
├── best_model_0001.h5
├── best_model_0002.h5
├── best_model_0003.h5
├── best_model_0005.h5
├── best_model_0006.h5
├── posenet_arrow.h5                # kerasモデル(HDF5形式)
├── posenet_arrow.keras             # kerasモデル
├── posenet_arrow_quantized.keras   # 量子化済みkerasモデル
└── posenet_arrow_quantized.tflite # 量子化済みtfliteモデル

Step 2: 学習済みモデルを量子化・変換する

次に量子化済みのモデルをAI Cameraにデプロイできる形にコンバートを行います。

ここの作業は引き続きPCのvenvの環境内で行ってください。
AIモデルのコンバートはモデルを作成したときと同じTensorFlowのバージョンで行う必要があります。

以下の公式ドキュメントに従って、進めていきます。

まずはedge-mdtライブラリをインストールします。
Edge-MDT (Model Development Toolkit)とは、IMX500向けにモデルを圧縮、量子化、コンバートするためのツール群になります。
それぞれPyTorch,TensorFlow向けがあります。今回はTensorFlow向けをインストールします。

pip install edge-mdt[tf]

次にモデルの量子化を行います。
-oはAIモデルを出力するために新規に作成するフォルダ名を指定しており、以下の例だとconvert_resultに出力されます。

cd ./samples/model/posenet_arrow/
imxconv-tf -i posenet_arrow_quantized.keras -o convert_result

上記のコマンドを実行すると変換が始まります。

およそ30秒程度でコンバートが完了します。
コンバートが完了すると以下のようにファイルが生成されます。

./samples/model/posenet_arrow/convert_result/
├── dnnParams.xml                           # ネットワークのパラメータ定義
├── packerOut.zip                           # ⭐️モデルをIMX向けにパッケージ化した成果物
├── posenet_arrow_quantized_MemoryReport.json # メモリ使用量レポート
└── posenet_arrow_quantized.pbtxt           # TensorFlow GraphDefテキスト表現

packerOut.zip が正しく生成されていればStep2は完了です。

Step 3: AIカメラ用にパッケージ化する

⭐️ここからはRaspberry Piでの操作になります⭐️

まず、先ほどのコンバートしたフォルダをRaspberry Piに移動します。
Step2で生成したpackerOut.zipをRaspberry Piの任意の場所に置いておきます。ファイルの移動にはscpやお好みのツールをご利用ください。

次に、RaspberrPiにてモデルのパッケージツールをインストールします。

sudo apt install imx500-tools

インストールが無事完了すると、次に変換を行います。
以下のコマンドでは、packerOut.zipを指定してパッケージ後の出力ファイルをrpk_output_folderに保存します。

imx500-package -i packerOut.zip -o rpk_output_folder

変換が完了すると.rpkの拡張子のパッケージ済みAIモデルが作成されます。

./rpk_output_folder/
└── network.rpk           # ⭐️AIカメラにデプロイする最終的な成果物

network.rpkのファイルが生成されていればStep３は完了です。

Step 4: Raspberry Pi上で実行する

Raspberry Pi上で必要なライブラリをインストールしていきます。

今回キーポイントの可視化をするにあたり、Raspberry Pi AI Camera向けのアプリ開発用ライブラリであるModlib(module-library)を使用します。
このライブラリは、検出モデルのサンプルをいくつか備えており、事前学習済みモデルからカスタムモデルまで簡単に実装することができます。
今回はAIモデルの可視化をなるべくシンプルな実装で実現するために、本ライブラリを使ってみます。

Modlib(module-library)をインストールします。

pip install git+https://github.com/SonySemiconductorSolutions/aitrios-rpi-application-module-library.git

次に可視化するためのスクリプトを準備します。

今回は矢印の7つの頂点を検出しそれぞれ色をつけて可視化するスクリプトを作成します。
必要に応じてAIモデル(rpk形式)のファイルのパスを変更ください。

可視化するスクリプト例

from pathlib import Path
from typing import List
import argparse

import cv2
import numpy as np

from modlib.devices import AiCamera
from modlib.models import COLOR_FORMAT, Model, MODEL_TYPE
from modlib.models.results import Poses
from modlib.models.post_processors import pp_personlab

# TODO: Update this path to point to your actual model file
MODEL_PACKED = Path("./path/to/your/model.rpk")


class ArrowPosenet(Model):

    def __init__(
        self,
        weights,
        model_type=MODEL_TYPE.RPK_PACKAGED,
        is_quantized: bool = True,
    ):
        super().__init__(
            model_file=weights,
            model_type=model_type,
            color_format=COLOR_FORMAT.RGB,
            preserve_aspect_ratio=False,
        )

        self.num_kp = 7
        self.peak_thresh = 0.1
        self.nms_thresh = 0.05
        self.kp_radius = 3

        self.in_width = 481
        self.in_height = 353
        self.is_quantized = is_quantized

    def pre_process(self, img):
        img_resized = cv2.resize(
            img,
            (self.in_width, self.in_height),
            interpolation=cv2.INTER_AREA,
        )
        in_tensor = img_resized / 255.0
        in_tensor = np.expand_dims(in_tensor, axis=0)
        return img_resized, in_tensor

    def post_process(self, output_tensors: List[np.ndarray]) -> Poses:
        edges = [
            (0, 1), (0, 6), (1, 2), (2, 3),
            (3, 4), (4, 5), (5, 6),
        ]
        
        return pp_personlab(
            output_tensors,
            num_keypoints=self.num_kp,
            edges=edges,
            peak_threshold=self.peak_thresh,
            nms_threshold=self.nms_thresh,
            kp_radius=self.kp_radius,
        )


def visualize_arrow(
    model_file,
    threshold=0.1,
    edge_thickness=4,
    draw_edges=False,
    draw_keypoints=True
):
    """Display arrow keypoints and skeleton edges on camera feed.
    
    Args:
        model_file: Path to the model file
        threshold: Detection threshold
        edge_thickness: Thickness of skeleton edges
        draw_edges: Whether to draw skeleton edges (default: False)
        draw_keypoints: Whether to draw keypoints (default: True)
    """

    model_path = Path(model_file)
    if not model_path.exists():
        raise FileNotFoundError(f"Model file not found: {model_path}")

    device = AiCamera()
    model = ArrowPosenet(weights=model_path)
    device.deploy(model)

    try:
        with device as stream:
            for frame in stream:
                if not hasattr(frame, "image") or frame.image is None:
                    continue

                if (
                    not hasattr(frame, "detections")
                    or frame.detections.n_detections == 0
                ):
                    frame.display()
                    continue

                try:
                    annotation_threshold = threshold
                    colors = [
                        (0, 0, 255),
                        (0, 255, 0),
                        (255, 0, 0),
                        (0, 255, 255),
                        (255, 0, 255),
                        (255, 255, 0),
                        (128, 128, 255),
                    ]
                    radius = 8
                    h, w, _ = frame.image.shape
                    poses = frame.detections
                    
                    keypoint_map = {}
                    
                    for i in range(poses.n_detections):
                        for k in range(model.num_kp):
                            score = poses.keypoint_scores[i, k]
                            if score > annotation_threshold:
                                x = int(poses.keypoints[i, k, 0] * w)
                                y = int(poses.keypoints[i, k, 1] * h)
                                if x > 0 and x < w and y > 0 and y < h:
                                    keypoint_map[k] = (x, y)
                    
                    if draw_edges:
                        edges = [
                            (0, 1), (0, 6), (1, 2), (2, 3),
                            (3, 4), (4, 5), (5, 6),
                        ]
                        edge_color = (255, 255, 255)
                        
                        for start_kp, end_kp in edges:
                            if (start_kp in keypoint_map and
                                    end_kp in keypoint_map):
                                x1, y1 = keypoint_map[start_kp]
                                x2, y2 = keypoint_map[end_kp]
                                cv2.line(frame.image, (x1, y1), (x2, y2),
                                         edge_color, edge_thickness)
                    
                    if draw_keypoints:
                        for k, (x, y) in keypoint_map.items():
                            col = colors[k % len(colors)]
                            cv2.circle(frame.image, (x, y), radius, col, -1)
                    
                    frame.display()
                except Exception:
                    continue

    except KeyboardInterrupt:
        pass


def main():
    parser = argparse.ArgumentParser(
        description="Arrow keypoint detection visualization"
    )
    parser.add_argument(
        "--model_file",
        type=str,
        default=str(MODEL_PACKED),
        help="Path to the model file"
    )
    parser.add_argument(
        "--threshold",
        type=float,
        default=0.1,
        help="Detection threshold (default: 0.1)"
    )
    parser.add_argument(
        "--edge_thickness",
        type=int,
        default=4,
        help="Thickness of skeleton edges (default: 4)"
    )
    parser.add_argument(
        "--draw_edges",
        action="store_true",
        default=False,
        help="Draw skeleton edges (default: False)"
    )
    args = parser.parse_args()
    
    # Keypoints are drawn by default
    draw_keypoints = True
    
    visualize_arrow(
        model_file=args.model_file,
        threshold=args.threshold,
        edge_thickness=args.edge_thickness,
        draw_edges=args.draw_edges,
        draw_keypoints=draw_keypoints
    )


if __name__ == "__main__":
    main()

上記のサンプルスクリプト例をもとに、visualize_arrow.pyを作成し実行します。

python3 visualize_arrow.py

実行結果

上記のプログラムを実行すると、自動的にAIモデルのデプロイが行われてWindowが表示されます。
実行結果はこのようになります。

正しくそれぞれの頂点を検出することができています。

応用：自作データセットに置き換えるには？

.ini ファイル内の以下の項目を書き換えることで、任意のデータで学習が可能です。

[DATASET] NAME = YourDatasetName
[MODEL] CLASS_NUM = キーポイント数
[TRAINER] CONFIG = あなたのYAML設定ファイル

データセットを置き換えるために必要の詳細は以下に記載しています。
https://github.com/SonySemiconductorSolutions/aitrios-rpi-training-samples/blob/main/docs/Posenet_NewDataset.md

困った時は

もし、記事の途中でうまくいかなかった場合は、気軽にこの記事にコメントいただければと思います。

もしRaspberryPiに関連する疑問がある場合は下記フォーラムもご確認、ご活用ください。

Raspberry Pi Forums

まとめ

Raspberry Pi AI カメラで利用できるキーポイント抽出のモデルの再学習ツールを紹介しました。
本記事では矢印の辺を検出するサンプルを紹介しました。今回はあくまでサンプルのため、矢印を検出するAIモデルを作成してみましたが、データセットを変更することでさまざまなエッジAIアプリケーションに応用可能です。

キーポイント検出モデルは人の骨格推定だけではなく、さまざまなケースケースに応用することもできるため、ぜひあなた自身のユースケースにもチャレンジしていただければと思います。

最後まで記事を読んでいただきありがとうございました。
もし記事が参考になりましたらいいねしていただければ幸いです。