RSDK-9495 - Add method which converts a pointcloud into an octree. (#…

…4608)
viamrobotics · Dec 9, 2024 · 8eebdce · 8eebdce
1 parent dd727dc
commit 8eebdce
Show file tree

Hide file tree

Showing 3 changed files with 73 additions and 27 deletions.
diff --git a/pointcloud/basic_octree_test.go b/pointcloud/basic_octree_test.go
@@ -633,24 +633,33 @@ func testPCDToBasicOctree(t *testing.T, artifactPath string) {
 	validateBasicOctree(t, basicOct, basicOct.center, basicOct.sideLength)
 }
 
-func TestCachedMaxProbability(t *testing.T) {
+func createPopulatedOctree(sign int) (*BasicOctree, error) {
 	center := r3.Vector{X: 0, Y: 0, Z: 0}
 	side := 2.0
+	octree, err := createNewOctree(center, side)
+	if err != nil {
+		return nil, err
+	}
+	pointsAndData := []PointAndData{
+		{P: r3.Vector{X: 0, Y: 0, Z: 0}, D: NewValueData(2 * sign)},
+		{P: r3.Vector{X: .5, Y: 0, Z: 0}, D: NewValueData(3 * sign)},
+		{P: r3.Vector{X: .5, Y: 0, Z: .5}, D: NewValueData(10 * sign)},
+		{P: r3.Vector{X: .5, Y: .5, Z: 0}, D: NewValueData(1 * sign)},
+		{P: r3.Vector{X: .55, Y: .55, Z: 0}, D: NewValueData(4 * sign)},
+		{P: r3.Vector{X: -.55, Y: -.55, Z: 0}, D: NewValueData(5 * sign)},
+		{P: r3.Vector{X: .755, Y: .755, Z: 0}, D: NewValueData(6 * sign)},
+	}
 
-	t.Run("get the max val from an octree", func(t *testing.T) {
-		octree, err := createNewOctree(center, side)
-		test.That(t, err, test.ShouldBeNil)
-		pointsAndData := []PointAndData{
-			{P: r3.Vector{X: 0, Y: 0, Z: 0}, D: NewValueData(2)},
-			{P: r3.Vector{X: .5, Y: 0, Z: 0}, D: NewValueData(3)},
-			{P: r3.Vector{X: .5, Y: 0, Z: .5}, D: NewValueData(10)},
-			{P: r3.Vector{X: .5, Y: .5, Z: 0}, D: NewValueData(1)},
-			{P: r3.Vector{X: .55, Y: .55, Z: 0}, D: NewValueData(4)},
-			{P: r3.Vector{X: -.55, Y: -.55, Z: 0}, D: NewValueData(5)},
-			{P: r3.Vector{X: .755, Y: .755, Z: 0}, D: NewValueData(6)},
-		}
+	err = addPoints(octree, pointsAndData)
+	if err != nil {
+		return nil, err
+	}
+	return octree, nil
+}
 
-		err = addPoints(octree, pointsAndData)
+func TestCachedMaxProbability(t *testing.T) {
+	t.Run("get the max val from an octree", func(t *testing.T) {
+		octree, err := createPopulatedOctree(1)
 		test.That(t, err, test.ShouldBeNil)
 
 		validateBasicOctree(t, octree, octree.center, octree.sideLength)
@@ -675,19 +684,7 @@ func TestCachedMaxProbability(t *testing.T) {
 	})
 
 	t.Run("setting negative values", func(t *testing.T) {
-		octree, err := createNewOctree(center, side)
-		test.That(t, err, test.ShouldBeNil)
-		pointsAndData := []PointAndData{
-			{P: r3.Vector{X: 0, Y: 0, Z: 0}, D: NewValueData(-2)},
-			{P: r3.Vector{X: .5, Y: 0, Z: 0}, D: NewValueData(-3)},
-			{P: r3.Vector{X: .5, Y: 0, Z: .5}, D: NewValueData(-10)},
-			{P: r3.Vector{X: .5, Y: .5, Z: 0}, D: NewValueData(-1)},
-			{P: r3.Vector{X: .55, Y: .55, Z: 0}, D: NewValueData(-4)},
-			{P: r3.Vector{X: -.55, Y: -.55, Z: 0}, D: NewValueData(-5)},
-			{P: r3.Vector{X: .755, Y: .755, Z: 0}, D: NewValueData(-6)},
-		}
-
-		err = addPoints(octree, pointsAndData)
+		octree, err := createPopulatedOctree(-1)
 		test.That(t, err, test.ShouldBeNil)
 
 		validateBasicOctree(t, octree, octree.center, octree.sideLength)

diff --git a/pointcloud/pointcloud_utils.go b/pointcloud/pointcloud_utils.go
@@ -91,3 +91,29 @@ func StatisticalOutlierFilter(meanK int, stdDevThresh float64) (func(PointCloud)
 	}
 	return filterFunc, nil
 }
+
+// ToBasicOctree takes a pointcloud object and converts it into a basic octree.
+func ToBasicOctree(cloud PointCloud) (*BasicOctree, error) {
+	if basicOctree, ok := cloud.(*BasicOctree); ok {
+		return basicOctree, nil
+	}
+
+	center := getCenterFromPcMetaData(cloud.MetaData())
+	maxSideLength := getMaxSideLengthFromPcMetaData(cloud.MetaData())
+	basicOctree, err := NewBasicOctree(center, maxSideLength)
+	if err != nil {
+		return nil, err
+	}
+	var iterateError error
+	cloud.Iterate(0, 0, func(p r3.Vector, d Data) bool {
+		if err = basicOctree.Set(p, d); err != nil {
+			iterateError = err
+			return false
+		}
+		return true
+	})
+	if iterateError != nil {
+		return nil, err
+	}
+	return basicOctree, nil
+}
diff --git a/pointcloud/pointcloud_utils_test.go b/pointcloud/pointcloud_utils_test.go
@@ -70,3 +70,26 @@ func TestPrune(t *testing.T) {
 	test.That(t, len(clouds), test.ShouldEqual, 1)
 	test.That(t, clouds[0].Size(), test.ShouldEqual, 5)
 }
+
+func TestToOctree(t *testing.T) {
+	pc := newBigPC()
+	tree, err := ToBasicOctree(pc)
+	test.That(t, err, test.ShouldBeNil)
+	pc.Iterate(0, 0, func(p r3.Vector, d Data) bool {
+		treeData, b := tree.At(p.X, p.Y, p.Z)
+		test.That(t, b, test.ShouldBeTrue)
+		test.ShouldResemble(t, treeData, d)
+		return true
+	})
+
+	basicTree, err := createPopulatedOctree(1)
+	test.That(t, err, test.ShouldBeNil)
+	tree, err = ToBasicOctree(basicTree)
+	test.That(t, err, test.ShouldBeNil)
+	basicTree.Iterate(0, 0, func(p r3.Vector, d Data) bool {
+		treeData, b := tree.At(p.X, p.Y, p.Z)
+		test.That(t, b, test.ShouldBeTrue)
+		test.ShouldResemble(t, treeData, d)
+		return true
+	})
+}