データフォーマット変換

次に学習に必要なデータセットを準備します。
学習に必要なデータは既にOmniverse Replicatorで作成していますが、出力された合成データをTAO Toolkitで学習可能なデータフォーマットに変換する必要があります。
Omniverse Replicatorで出力されたzipファイルを解凍するとpng, npy, jsonの3種類のデータが格納されていることがわかります。
それぞれ以下の内容のファイルになっています。

・png: 合成されたRGB画像データ
・npy: 合成画像内アセットの位置を示すBounding Boxデータ
・json: 合成画像内アセットの内容を示すClassデータ


これらのデータをTAO Toolkitで学習するためにKITTIフォーマットというデータフォーマットに変換します。
KITTIフォーマットとはClassデータ、Bounding Boxデータを含むアノテーション情報を、以下のように画像ごとのtxtファイルで出力するデータフォーマットのことです。

それではnpyファイルおよびjsonファイルから必要な情報を抽出してKITTIフォーマットに変換していきます。
変換にはPythonスクリプトを使用します。

以上の処理が完了すると以下のようにデータセットが整理されます。

整理されたデータの数を確認してみます。

このままでも学習の実行は可能ですが、少々データ数が少ないため、次にデータ拡張（Data Augmentation）という処理をおこなっていきます。

データ拡張

データ拡張とは、学習データの環境条件を変更することによって学習データを水増しする手法です。
データ拡張は学習データ数が十分に確保できていない場合に有効な手法で、Data Augmentationとも呼ばれます。
このデータ拡張は、Pythonのライブラリ等を駆使して行われるのが一般的ですが、TAO Toolkitにはデータ拡張をおこなう機能が備わっているため、学習と同様にコマンド1つで簡単にデータ拡張を実行できます。

NVIDIA TAO Toolkitによるデータ拡張は以下のspecファイルの内容に基づいておこなわれます。

この内容では画像の90度回転、0～20度の色相回転、画像ぼかし、輝度変更が実行されます。
以下のコマンドでデータ拡張がおこなわれます。

データ拡張を行った後の画像とBounding Boxを確認してみます。

元の画像と似てはいますが、異なる画像になっていることが確認できます。
このようにデータ拡張をおこなうことで、学習に使用するデータを水増しすることができます。

以上でデータセットの準備が完了しました。

転移学習の実行

これまでに準備してきたデータセット、事前学習済みモデルを使用してTAO Toolkitによる転移学習をおこないます。
以下のspecファイルの内容に基づいて転移学習が実行されます。エポック数やバッチサイズ等の学習に関わる設定はこのspecファイルを編集することで変更できます。

以下のコマンドを実行すると転移学習が実行されます。

学習が完了すると、学習済みモデルが出力されます。

推論の実行

学習が完了したら推論結果を確認してみます。

AIモデルが期待通りに人形を検出できているため、学習を終了します。