Potree to PLY conversion utility

examples/potree-utils/extract_pointcloud.go

@@ -1,18 +1,157 @@
 package main
 
 import (
+	"bufio"
+	"encoding/binary"
 	"fmt"
 	"io"
+	"log"
+	"math"
 	"os"
+	"runtime"
+	"strings"
+	"sync"
+	"time"
 
 	"github.com/EliCDavis/polyform/formats/ply"
 	"github.com/EliCDavis/polyform/formats/potree"
 	"github.com/EliCDavis/polyform/modeling"
 	"github.com/urfave/cli/v2"
 )
 
-func buildModel(octreeFile *os.File, node *potree.OctreeNode, metadata *potree.Metadata, includeChildren bool) (*modeling.Mesh, error) {
-	_, err := octreeFile.Seek(int64(node.ByteOffset), 0)
+func buildModelWorker(
+	ctx *cli.Context,
+	nodes <-chan *potree.OctreeNode,
+	metadata *potree.Metadata,
+	largestPointCount int,
+	plyOut io.Writer,
+	mutex *sync.Mutex,
+	out chan<- int,
+) {
+	octreeFile, err := openOctreeFile(ctx)
+	if err != nil {
+		panic(err)
+	}
+	defer octreeFile.Close()
+	pointsProcessed := 0
+	potreeBuf := make([]byte, largestPointCount*metadata.BytesPerPoint())
+	plyBuf := make([]byte, 27)
+	for node := range nodes {
+		_, err := octreeFile.Seek(int64(node.ByteOffset), 0)
+		if err != nil {
+			panic(err)
+		}
+
+		_, err = io.ReadFull(octreeFile, potreeBuf[:node.ByteSize])
+		if err != nil {
+			panic(err)
+		}
+
+		cloud := potree.LoadNode(*node, *metadata, potreeBuf[:node.ByteSize])
+		count := cloud.PrimitiveCount()
+		if count == 0 {
+			continue
+		}
+		positions := cloud.Float3Attribute(modeling.PositionAttribute)
+		colors := cloud.Float3Attribute(modeling.ColorAttribute)
+
+		mutex.Lock()
+
+		endien := binary.LittleEndian
+		for i := 0; i < count; i++ {
+			p := positions.At(i)
+			bits := math.Float64bits(p.X())
+			endien.PutUint64(plyBuf, bits)
+			bits = math.Float64bits(p.Y())
+			endien.PutUint64(plyBuf[8:], bits)
+			bits = math.Float64bits(p.Z())
+			endien.PutUint64(plyBuf[16:], bits)
+
+			c := colors.At(i).Scale(255)
+			plyBuf[24] = byte(c.X())
+			plyBuf[25] = byte(c.Y())
+			plyBuf[26] = byte(c.Z())
+
+			_, err = plyOut.Write(plyBuf)
+			if err != nil {
+				panic(err)
+			}
+		}
+
+		mutex.Unlock()
+
+		pointsProcessed += count
+	}
+
+	out <- pointsProcessed
+}
+
+func buildModelWithChildren(ctx *cli.Context, root *potree.OctreeNode, metadata *potree.Metadata) error {
+
+	f, err := os.Create(ctx.String("out"))
+	if err != nil {
+		return err
+	}
+	defer f.Close()
+
+	largestPointCount := 0
+	root.Walk(func(o *potree.OctreeNode) {
+		if o.NumPoints > uint32(largestPointCount) {
+			largestPointCount = int(o.NumPoints)
+		}
+	})
+
+	header := ply.Header{
+		Format: ply.BinaryLittleEndian,
+		Elements: []ply.Element{
+			{
+				Name:  ply.VertexElementName,
+				Count: int64(root.PointCount()),
+				Properties: []ply.Property{
+					&ply.ScalarProperty{PropertyName: "x", Type: ply.Double},
+					&ply.ScalarProperty{PropertyName: "y", Type: ply.Double},
+					&ply.ScalarProperty{PropertyName: "z", Type: ply.Double},
+					&ply.ScalarProperty{PropertyName: "red", Type: ply.UChar},
+					&ply.ScalarProperty{PropertyName: "green", Type: ply.UChar},
+					&ply.ScalarProperty{PropertyName: "blue", Type: ply.UChar},
+				},
+			},
+		},
+	}
+
+	bufWriter := bufio.NewWriter(f)
+	err = header.Write(bufWriter)
+	if err != nil {
+		return err
+	}
+
+	workerCount := runtime.NumCPU()
+	jobs := make(chan *potree.OctreeNode, workerCount)
+	meshes := make(chan int, workerCount)
+
+	mutex := &sync.Mutex{}
+	for i := 0; i < workerCount; i++ {
+		go buildModelWorker(ctx, jobs, metadata, largestPointCount, bufWriter, mutex, meshes)
+	}
+
+	root.Walk(func(o *potree.OctreeNode) { jobs <- o })
+	close(jobs)
+
+	for i := 0; i < workerCount; i++ {
+		<-meshes
+	}
+
+	return bufWriter.Flush()
+}
+
+func buildModel(ctx *cli.Context, node *potree.OctreeNode, metadata *potree.Metadata) (*modeling.Mesh, error) {
+	octreeFile, err := openOctreeFile(ctx)
+	if err != nil {
+		return nil, err
+	}
+	defer octreeFile.Close()
+
+	_, err = octreeFile.Seek(int64(node.ByteOffset), 0)
 	if err != nil {
 		return nil, err
 	}
@@ -24,16 +163,22 @@ func buildModel(octreeFile *os.File, node *potree.OctreeNode, metadata *potree.M
 	}
 
 	mesh := potree.LoadNode(*node, *metadata, buf)
-	if includeChildren {
-		for _, c := range node.Children {
-			childMesh, err := buildModel(octreeFile, c, metadata, true)
-			if err != nil {
-				return nil, err
-			}
-			mesh = mesh.Append(*childMesh)
+	return &mesh, nil
+}
+
+func findNode(name string, node *potree.OctreeNode) (*potree.OctreeNode, error) {
+	if name == node.Name {
+		return node, nil
+	}
+
+	for _, c := range node.Children {
+		if strings.Index(name, c.Name) == 0 {
+			return findNode(name, c)
 		}
+		log.Println(c.Name)
 	}
-	return &mesh, nil
+
+	return nil, fmt.Errorf("%s can't find child node with name %s", node.Name, name)
 }
 
 var ExtractPointcloudCommand = &cli.Command{
@@ -64,19 +209,23 @@ var ExtractPointcloudCommand = &cli.Command{
 			return err
 		}
 
-		octreeFile, err := openOctreeFile(ctx)
+		startNode, err := findNode(ctx.String("node"), hierarchy)
 		if err != nil {
 			return err
 		}
-		defer octreeFile.Close()
 
-		mesh, err := buildModel(octreeFile, hierarchy, metadata, ctx.Bool("include-children"))
-		if err != nil {
+		if ctx.Bool("include-children") {
+			start := time.Now()
+			err = buildModelWithChildren(ctx, startNode, metadata)
+			log.Printf("PLY written in %s", time.Since(start))
 			return err
 		}
 
+		mesh, err := buildModel(ctx, startNode, metadata)
+		if err != nil {
+			return err
+		}
 		fmt.Fprintf(ctx.App.Writer, "Writing pointcloud with %d points to %s", mesh.Indices().Len(), ctx.String("out"))
-
 		return ply.SaveBinary(ctx.String("out"), *mesh)
 	},
 }

examples/potree-utils/util.go

@@ -48,5 +48,5 @@ func openOctreeFile(ctx *cli.Context) (*os.File, error) {
 		octreePath = filepath.Join(filepath.Dir(metadataPath), "octree.bin")
 	}
 
-	return os.Open(octreePath)
+	return os.OpenFile(octreePath, os.O_RDONLY, 0)
 }

formats/ply/ascii_reader.go

@@ -152,7 +152,7 @@ func (ar *AsciiReader) readVertexData(element Element, approvedData map[string]b
 		}
 	}
 
-	i := 0
+	var i int64 = 0
 	for i < element.Count {
 		ar.scanner.Scan()
 
@@ -292,7 +292,7 @@ func (ar *AsciiReader) readFaceData(element Element) ([]int, []vector2.Float64,
 		}
 	}
 
-	for i := 0; i < element.Count; i++ {
+	for i := int64(0); i < element.Count; i++ {
 		ar.scanner.Scan()
 		contents := strings.Fields(ar.scanner.Text())
 		offset := 0
@@ -313,7 +313,7 @@ func (ar *AsciiReader) ReadMesh(vertexAttributes map[string]bool) (*modeling.Mes
 	var triData []int
 	var uvData []vector2.Float64
 	for _, element := range ar.elements {
-		if element.Name == "vertex" {
+		if element.Name == VertexElementName {
 			data, err := ar.readVertexData(element, vertexAttributes)
 			if err != nil {
 				return nil, err

formats/ply/binary_reader.go

@@ -259,7 +259,7 @@ func (le *BinaryReader) readVertexData(element Element, approvedData map[string]
 		}
 	}
 
-	i := 0
+	var i int64 = 0
 	buf := make([]byte, nextOffset)
 	for i < element.Count {
 		_, err := io.ReadFull(le.reader, buf)
@@ -462,7 +462,7 @@ func (le *BinaryReader) readFaceData(element Element) ([]int, []vector2.Float64,
 		}
 	}
 
-	for i := 0; i < element.Count; i++ {
+	for i := int64(0); i < element.Count; i++ {
 		for _, reader := range readers {
 			reader(le.reader)
 		}
@@ -476,7 +476,7 @@ func (le *BinaryReader) ReadMesh(vertexAttributes map[string]bool) (*modeling.Me
 	var triData []int
 	var uvData []vector2.Float64
 	for _, element := range le.elements {
-		if element.Name == "vertex" {
+		if element.Name == VertexElementName {
 			data, err := le.readVertexData(element, vertexAttributes)
 			if err != nil {
 				return nil, err

formats/ply/element.go

@@ -7,7 +7,7 @@ import (
 
 type Element struct {
 	Name       string     `json:"name"`
-	Count      int        `json:"count"`
+	Count      int64      `json:"count"`
 	Properties []Property `json:"properties"`
 }
 

formats/ply/format.go

@@ -7,3 +7,5 @@ const (
 	BinaryBigEndian
 	BinaryLittleEndian
 )
+
+const VertexElementName = "vertex"

formats/ply/reader.go

@@ -196,7 +196,7 @@ func ReadHeader(in io.Reader) (Header, error) {
 				return header, errors.New("illegal element line in ply header")
 			}
 
-			elementCount, err := strconv.Atoi(contents[2])
+			elementCount, err := strconv.ParseInt(contents[2], 10, 64)
 			if err != nil {
 				return header, fmt.Errorf("unable to parse element count: %w", err)
 			}
@@ -258,9 +258,9 @@ func ReadMesh(in io.Reader) (*modeling.Mesh, error) {
 	return mesh, nil
 }
 
-func buildVertexElements(attributes []string, size int) Element {
+func buildVertexElements(attributes []string, size int64) Element {
 	vertexElement := Element{
-		Name:       "vertex",
+		Name:       VertexElementName,
 		Count:      size,
 		Properties: make([]Property, 0),
 	}

formats/ply/write.go

@@ -17,7 +17,7 @@ func BuildHeaderFromModel(model modeling.Mesh, format Format) Header {
 	header := Header{
 		Format: format,
 		Elements: []Element{
-			buildVertexElements(model.Float3Attributes(), model.AttributeLength()),
+			buildVertexElements(model.Float3Attributes(), int64(model.AttributeLength())),
 		},
 	}
 
@@ -49,7 +49,7 @@ func BuildHeaderFromModel(model modeling.Mesh, format Format) Header {
 
 		header.Elements = append(header.Elements, Element{
 			Name:       "face",
-			Count:      model.PrimitiveCount(),
+			Count:      int64(model.PrimitiveCount()),
 			Properties: faceProperties,
 		})
 	}

formats/potree/octree_node.go

@@ -26,6 +26,13 @@ type OctreeNode struct {
 	HierarchyByteSize   uint64
 }
 
+func (on *OctreeNode) Walk(f func(o *OctreeNode)) {
+	f(on)
+	for _, c := range on.Children {
+		c.Walk(f)
+	}
+}
+
 func (on OctreeNode) Height() int {
 	if len(on.Children) == 0 {
 		return 0