Merge pull request #16 from HiraiKyo/feature/clustering

平面上の点群のDBSCANクラスタリング実装

README.md

@@ -8,6 +8,18 @@ Basic libraries for manipulating point cloud.
 pip install git+https://github.com/HiraiKyo/ply-processor-basics
 ```
 
+## Development
+
+### Running test
+
+`visual`タグはopen3d.geometry等で表示を確認する用なので、テスト実行時は外す
+
+```sh
+poetry run pytest -s {filepath} -m "not visual"
+```
+
+TDD開発時にopen3dで表示を確認しつつ進める場合には、そのテストに`@pytest.mark.visual`タグを付けて自動テストに影響しないようにする。
+
 ## Methods
 
 ### STL
@@ -34,6 +46,8 @@ pip install git+https://github.com/HiraiKyo/ply-processor-basics
 
 #### `points.clip_by_plane`
 
+#### `points.plane_clustering`
+
 #### `points.ransac.detect_plane`
 
 #### `points.ransac.detect_circle`

ply_processor_basics/points/__init__.py

@@ -1,4 +1,5 @@
 from .clip_by_plane import clip_by_plane as clip_by_plane
+from .clustering import plane_clustering as plane_clustering
 from .get_distances_to_line import get_distances_to_line as get_distances_to_line
 from .get_distances_to_plane import get_distances_to_plane as get_distances_to_plane
 from .rotate_euler import rotate_euler as rotate_euler

ply_processor_basics/points/clustering.py

@@ -0,0 +1,20 @@
+import numpy as np
+from numpy.typing import NDArray
+from sklearn.cluster import DBSCAN
+
+
+def plane_clustering(points: NDArray[np.floating], eps: float = 1.0, min_samples: int = 100):
+    """
+    DBSCANによる平面上の点群のクラスタリングを行う
+
+    :param points: 平面上の点群(N, 2)
+    :return: クラスタ点数の多い順にソートされたクラスタ点群ポインタ(N, M)
+    """
+    dbscan = DBSCAN(eps=eps, min_samples=min_samples, metric="euclidean")
+    clusters = dbscan.fit_predict(points[:, :2])
+    unique_clusters, counts = np.unique(clusters, return_counts=True)
+    # クラスタリングできなかった点群(-1)は除外
+    unique_clusters = unique_clusters[unique_clusters != -1]
+    # 各クラスタの点群ポインタを返す
+    cluster_indices = [np.where(clusters == cluster)[0] for cluster in unique_clusters]
+    return cluster_indices

tests/points/ransac/test_detect_plane.py

@@ -12,7 +12,7 @@ def test_success():
     assert model is not None
     assert len(inliers) > 10000
     assert len(model) == 4
-    # およそZ軸方向に平面の法線ベクトルが向いていることを確認
+    # サンプルデータはZ軸方向に平面の法線ベクトルが向いている
     normalized = model[:3] / np.linalg.norm(model[:3])
     # normalizedが[0, 0, 1]もしくは[0, 0, -1]に近いことを確認
     assert np.allclose(normalized, [0, 0, 1], atol=0.1) or np.allclose(normalized, [0, 0, -1], atol=0.1)

tests/points/test_clustering.py

@@ -0,0 +1,36 @@
+import numpy as np
+import open3d as o3d
+import pytest
+
+from ply_processor_basics.points import plane_clustering
+from ply_processor_basics.points.ransac import detect_plane
+
+
+@pytest.mark.parametrize("plypath", ["data/samples/sample_clustering.ply"])
+def test_success(plypath):
+    pcd = o3d.io.read_point_cloud(plypath)
+    points = np.asarray(pcd.points)
+
+    inliers, plane_model = detect_plane(points, threshold=1.0)
+    # EPS=10.0でサンプルデータ点群のクラスタリングが正常動作した事を確認
+    clusters_indices = plane_clustering(points[inliers], eps=10.0)
+    # クラスタが2つ認識されることを期待
+    assert len(clusters_indices) == 2
+
+
+@pytest.mark.parametrize("plypath", ["data/samples/sample_clustering.ply"])
+@pytest.mark.visual
+def test_visualize(plypath):
+    pcd = o3d.io.read_point_cloud(plypath)
+    points = np.asarray(pcd.points)
+
+    inliers, plane_model = detect_plane(points, threshold=1.0)
+    # サンプルデータはZ軸方向に平面の法線が存在するので座標変換は行わない
+    clusters_indices = plane_clustering(points[inliers], eps=10.0)  # EPS=10.0mmくらいで適切な
+    pcds = []
+    for cluster in clusters_indices:
+        pcd = o3d.geometry.PointCloud()
+        pcd.points = o3d.utility.Vector3dVector(points[inliers][cluster])
+        pcd.paint_uniform_color([np.random.rand(), np.random.rand(), np.random.rand()])
+        pcds.append(pcd)
+    o3d.visualization.draw_geometries(pcds)