Zenn
💬

Kubricで機械学習用の合成データを作成する

2022/07/03に公開
機械学習
合成データ
tech

※本記事は 2022年07月執筆時の情報です。

はじめに

機械学習プロジェクトでは良質な教師データが重要とされています。しかし、データ収集やアノテーション作業には手間と時間がかかります。そこでシミュレーションで生成される合成データの利用が提案されています。

Kubric は Google Research が公開している、機械学習用の合成データを生成するためのパイプラインです。3D モデルを使ってインスタンスセグメンテーション、深度画像、オプティカルフローなどのデータセットを作成することができます。物理シミュレーションには PyBullet、レンダリングには Blender が使われています。

Kubric のインストール

README と Installing を参考に Kubric が動く環境を作ります。

  • Kubric のリポジトリをクローンします。

    $ git clone https://github.com/google-research/kubric.git

  • 提供されている Docker イメージをプルしてサンプルを実行します。

    $ cd kubric
$ docker pull kubricdockerhub/kubruntu
$ docker run --rm --interactive \
           --user $(id -u):$(id -g) \
           --volume "$(pwd):/kubric" \
           kubricdockerhub/kubruntu \
           /usr/bin/python3 examples/helloworld.py
    • output ディレクトリが作られ画像が出力されます。

  • なお、自分で Docker イメージをビルドする場合は、wheel パッケージを先にビルドする必要があるようです。

    $ cd kubric
$ python3 setup.py bdist_wheel 
$ docker build -t kubricdockerhub/kubruntu -f docker/Kubruntu.Dockerfile .

Kubric のチュートリアル

Hello World

helloworld.py と Helloworld を参考に、上記で実行したサンプルを読んでいきます。

  • 以下でシーンとレンダリングの設定します。解像度 256x256 が設定されています。

    import kubric as kb
from kubric.renderer.blender import Blender as KubricRenderer

scene = kb.Scene(resolution=(256, 256))
renderer = KubricRenderer(scene)

  • 以下でシーンに床と球を追加します。

    scene += kb.Cube(scale=(10, 10, 0.1), position=(0, 0, -0.1))
scene += kb.Sphere(scale=1, position=(0, 0, 1.))

  • 以下でシーンに照明とカメラを追加します。

    scene += kb.DirectionalLight(position=(-1, -0.5, 3), look_at=(0, 0, 0), intensity=1.5)
scene += kb.PerspectiveCamera(position=(3, -0.5, 4), look_at=(0, 0, 1))

  • 以下でシーンをエクスポートできます。ローカル環境に Blender がインストールしてあればファイルを開いてシーンを確認することができます。

    renderer.save_state("helloworld.blend")

  • 以下でレンダリングを実行します。

    frame = renderer.render_still()

  • 以下でレンダリングされた映像を画像ファイルに保存します。

    • write_png() はカラー画像を通常の PNG として保存します。
    • write_palette_png() はセグメンテーション画像を保存しますが、カラーパレット（インデックスカラー）を使って個々の領域を見分けやすくしています。
    • write_scaled_png() は深度画像を 16 bit 画像として保存しているようです。暗い領域ほど浅い深度、明るい領域ほど深い深度を表しています。
    kb.write_png(frame["rgba"], "output/helloworld.png")
kb.write_palette_png(frame["segmentation"], "output/helloworld_segmentation.png")
kb.write_scaled_png(frame["depth"], "output/helloworld_depth.png")



物理シミュレーション

simulator.py と Simulator を参考に、物理シミュレーションのサンプルを読んでいきます。

  • 以下でシーンとレンダリングに加えて物理シミュレーションの設定をしています。

    scene = kb.Scene(resolution=(256, 256))
scene.frame_end = 48   # レンダリングするフレーム数
scene.frame_rate = 24  # フレームレート
scene.step_rate = 240  # 物理シミュレーションのフレームレート
simulator = KubricSimulator(scene)
renderer = KubricRenderer(scene, scratch_dir="./output")
    • 物理シミュレーションのフレームレートが 240Hz に設定されていますが、これは PyBullet のデフォルト値で、変更はあまり推奨されないようです。参考: PyBullet Quickstart Guide #setTimeStep

  • 以下でシーンに床を追加しています。static=True によって空間に固定されます。

    scene += kb.Cube(scale=(3, 3, 0.1), position=(0, 0, -0.1), static=True)

  • 以下で複数の球を追加します。

    spawn_region = [[-1, -1, 0], [1, 1, 1]]
rng = np.random.default_rng()
for i in range(8):
  position = rng.uniform(*spawn_region)
  velocity = rng.uniform([-1, -1, 0], [1, 1, 0])
  material = kb.PrincipledBSDFMaterial(color=kb.random_hue_color(rng=rng))
  sphere = kb.Sphere(scale=0.1, position=position, velocity=velocity, material=material)
  scene += sphere
  kb.move_until_no_overlap(sphere, simulator, spawn_region=spawn_region)
    • rng = np.random.default_rng() はNumpyの乱数ジェネレータです。rng.uniform() で一様分布を作ることができます。また kb.random_hue_color() の引数にも使われます。
    • kb.PrincipledBSDFMaterial() はパラメータによって幅広いマテリアルを表現できるそうです。この例ではランダムな色に設定されています。参考: Principled (プリンシプル) BSDF
    • kb.move_until_no_overlap() は球が干渉しないように位置を調整します。

  • 以下でシミュレーションを実行します。シミュレーション結果は renderer 内に保存されます。

    simulator.run()

  • 以下でシミュレーション結果をレンダリングします。連番のファイル名で出力されます。

    frames_dict = renderer.render()
kb.write_image_dict(frames_dict, "output")

  • ImageMagick を使うと連番画像から動画を生成できます。

    $ sudo apt install imagemagik
$ convert -delay 8 -loop 0 output/rgba_*.png output/simulator.gif

アセット

  • 以下のアセットの3D モデルを利用することができます。

  • 利用方法は執筆時点ではまだドキュメント化されていないのですが、challenges/movi/movi_c_worker.py を参考にすると、次のような形になると思います。

    kubasic = kb.AssetSource.from_manifest(
    "gs://kubric-public/assets/KuBasic/KuBasic.json")
material = kb.PrincipledBSDFMaterial()
teapot = kubasic.create(
    asset_id="teapot", material=material, position=(0, 0, 0.5))
scene += teapot
    • 参考にした movi_c_worker.py では3D モデルにランダムなテクスチャを設定していましたが、ここでは省略して単色にしています。

おわりに

Kubric の基本的な使い方を試してみました。上手く使いこなせばオリジナルのデータセットが簡単に作れそうです。フォトグラメトリ等で自作した 3D モデルを取り込んだりといった活用もやってみたいと思います。

Discussion