VoteNetを使った3D点群データの物体認識がすごそうなので早速動かしてみた

はじめに

VoteNet は　3D点群の物体検出をするすごいAIフレームワークです。
従来までは、3Dを2Dに変換してCNNで検出したりとかをしていたと思うのですが、VoteNetでは3Dのまま3Dバウンディングボックスで物体検出します。

詳しい説明は、下記リンクを参照してください。とてもわかりやすいです。
https://ai-scholar.tech/treatise/deephaf-ai-129/

この投稿では、実際に動かしてみるところに重きをおいています。

今回動作対象のものがある場所

下記リンクです。facebookresearchさんのGitHubです。
https://github.com/facebookresearch/votenet

環境

docker上に環境を構築しますが、そのホストの環境です。
ubuntu 18.04
NVIDIA GeForce RTX2080

インストール

docker環境の構築

別にdockerでなくてもいいのですが、GPU版のpytorchを使うのでCUDAやらcuDNNやらのバージョンがなんちゃらかんちゃらなので、dockerで作成します。

dockerは、GPUを使用するので、nvidia-docker-toolkitを使っています。下記のURLを参考にインストール済みです。
https://github.com/NVIDIA/nvidia-docker

GPUを使いたいので、pytorchが対応するCUDAやらのバージョンを確認します。
https://pytorch.org/
上記URLを参照すると、pytorchとその対応するCUDAやpythonのバージョンの組み合わせとインストール用のコマンドが確認できます。

CUDAは10.1を入れればいいようなので、下記コマンドでプルします。

# docker pull nvidia/cuda:10.1-cudnn7-devel-ubuntu18.04

プルが終わったら、起動します。

# docker run --gpus all -it --name pytorchgpu -p 8888:8888 -v ~/workspace/jupyter:/jupyter nvidia/cuda:10.1-cudnn7-devel-ubuntu18.04

GPUを使用するので、--gpus allをつけて起動します。-p 8888:8888 はjupyterの実行を見越しています。

pythonのインストール

出来上がった環境でPythonのバージョンを確認すると、

# python --version
bash: python: command not found

ってなって、pythonがデフォでインストールされていないようです。
ということで、pythonをインストールしましょう。

pythonは3.7をインストールします。venvを使用するので、それもインストールします。

# apt update
# apt install python3
# apt install python3-venv
# apt install python3.7
# apt install python3.7-venv
# apt install python3.7-dev

上記コマンドを実行しました。
後ほどのために、python3.7-devもインストールしておきましょう。
インストールが終わったらvenvで仮想の環境を作り、そこにpytorchをインストールします。

# cd ~
# mkdir env
# cd env
# python3.7 -m venv pytorch

作った仮想の環境にはいります。

# source pytorch/bin/activate
# python --version
Python 3.7.5

pytorchをインストールします。

# pip install torch

VoteNetインストール

gitコマンドでソースコードを取得します。

# git clone https://github.com/facebookresearch/votenet.git

必要なライブラリをインストールしておきます。numpyやplyfileなどが必要です。
trimesh は <2.35.40 である必要があるので、バージョンを指定しています。

# pip install numpy
# pip install wheel
# pip install plyfile
# pip install trimesh==2.35.39
# pip install matplotlib
# pip install opencv-python

opencvをインストールします。

# apt install libopencv-dev
# pip install opencv-python

セットアップします

# cd votenet/pointnet2/
# python setup.py install

セットアップに成功すると、下記のようなメッセージが表示されます。

Installed /root/env/pytorch/lib/python3.7/site-packages/pointnet2-0.0.0-py3.7-linux-x86_64.egg
Processing dependencies for pointnet2==0.0.0
Finished processing dependencies for pointnet2==0.0.0

デモンストレーション

デモンストレーション用のファイルを下記からダウンロードします。

votenet_pretrained_models.zipというファイルが、取得できます。
解凍して、votenetフォルダに配置します。unzipコマンドは適宜aptでインストールしてください。

# cd ..
# unzip [ZIP配置ディレクトリ]/votenet_pretrained_models.zip  -d .

デモンストレーションを実行します。

# python demo.py

以下のエラーが発生しました。

Traceback (most recent call last):
  File "demo.py", line 101, in <module>
    MODEL.dump_results(end_points, dump_dir, DC, True)
  File "/root/work/votenet/models/dump_helper.py", line 86, in dump_results
    pc_util.write_oriented_bbox(obbs[objectness_prob>DUMP_CONF_THRESH,:], os.path.join(dump_dir, '%06d_pred_confident_bbox.ply'%(idx_beg+i)))
  File "/root/work/votenet/utils/pc_util.py", line 421, in write_oriented_bbox
    mesh_list = trimesh.util.concatenate(scene.dump())
  File "/root/env/pytorch/lib/python3.7/site-packages/trimesh/scene/scene.py", line 499, in dump
    for node_name in self.graph.nodes_geometry:
  File "/root/env/pytorch/lib/python3.7/site-packages/trimesh/caching.py", line 88, in get_cached
    value = function(*args, **kwargs)
  File "/root/env/pytorch/lib/python3.7/site-packages/trimesh/scene/transforms.py", line 224, in nodes_geometry
    'geometry' in self.transforms.node[n]):
AttributeError: 'EnforcedForest' object has no attribute 'node'

このエラーは、networkxのバージョンが2.3未満でないと出力されるようですので、一旦削除して、バージョン指定してインストールしなおします。

# python
>>> import networkx
>>> print(networkx.__version__)
2.4

で、バージョン確認後アンインストールしt、下記のようにインストールします。

# pip uninstall networkx
# pip install networkx==2.2

もう一度実行してみましょう。

# python demo.py
Constructed model.
Loaded checkpoint /root/work/votenet/demo_files/pretrained_votenet_on_sunrgbd.tar (epoch: 180)
Loaded point cloud data: /root/work/votenet/demo_files/input_pc_sunrgbd.ply
Inference time: 0.078056
Finished detection. 9 object detected.
Dumped detection results to folder /root/work/votenet/demo_files/sunrgbd_results

今度は成功しました。demo_files/sunrgbd_resultsに結果が出力されます。

# ls -al demo_files/sunrgbd_results/
total 1512
drwxr-xr-x 2 root root    4096 Jan 28 13:12 .
drwxrwxr-x 3 root root    4096 Jan 28 11:58 ..
-rw-r--r-- 1 root root   16077 Jan 28 13:12 000000_aggregated_vote_pc.ply
-rw-r--r-- 1 root root    7057 Jan 28 13:12 000000_confident_proposal_pc.ply
-rw-r--r-- 1 root root 1253931 Jan 28 13:12 000000_pc.ply
-rw-r--r-- 1 root root   64721 Jan 28 13:12 000000_pred_bbox.ply
-rw-r--r-- 1 root root   28432 Jan 28 13:12 000000_pred_confident_bbox.ply
-rw-r--r-- 1 root root    2474 Jan 28 13:12 000000_pred_confident_nms_bbox.ply
-rw-r--r-- 1 root root   11799 Jan 28 13:12 000000_pred_nms_bbox.ply
-rw-r--r-- 1 root root   15995 Jan 28 13:12 000000_proposal_pc.ply
-rw-r--r-- 1 root root   64226 Jan 28 13:12 000000_seed_pc.ply
-rw-r--r-- 1 root root   63991 Jan 28 13:12 000000_vgen_pc.ply

この中で、000000_pred_confident_nms_bbox.plyが分類結果のバウンディングボックスです。

結果をCloudCompareで表示してみます。
認識元のファイルは、demo_filesフォルダの直下にあるinput_pc_sunrgbd.plyなので、これと、結果のファイルの二つを表示します。
※CloudCompareは3Dが表示できるフリーのツールです。

バウンディングボックスが塗りつぶされているので、みにくいです。DB TreeのMeshを選択して、Propertiesに表示されたMeshのWireframeのチェックをつけると、以下のような表示になります。

椅子と机が認識されています。

大成功です！

つぎは

学習させたいけど、アノテーションが大変そうです・・。
がんばります。