diff --git a/command_test.go b/command_test.go index a659d09..759d0a7 100644 --- a/command_test.go +++ b/command_test.go @@ -318,7 +318,8 @@ func TestAddSurface(t *testing.T) { }) } -func TestRelabelPointsInLabelRange(t *testing.T) { +func newLabeledTestCloud(t *testing.T) *pc.PointCloud { + t.Helper() header := pc.PointCloudHeader{ Fields: []string{"x", "y", "z", "label"}, Size: []int{4, 4, 4, 4}, @@ -356,7 +357,10 @@ func TestRelabelPointsInLabelRange(t *testing.T) { lt.Incr() lt.SetUint32(3) lt.Incr() + return pp +} +func TestRelabelPointsInLabelRange(t *testing.T) { c := newCommandContext(&dummyPCDIO{}, nil) testCases := map[string]struct { @@ -380,13 +384,17 @@ func TestRelabelPointsInLabelRange(t *testing.T) { for name, tt := range testCases { tt := tt t.Run(name, func(t *testing.T) { - c.SetPointCloud(pp, cloudMain) + // Edits mutate the cloud in place; use a fresh one per case. + c.SetPointCloud(newLabeledTestCloud(t), cloudMain) if err := c.RelabelPointsInLabelRange(tt.minLabel, tt.maxLabel, tt.newLabel); err != nil { t.Fatal(err) } - lt, err = c.editor.pp.Uint32Iterator("label") + lt, err := c.editor.pp.Uint32Iterator("label") + if err != nil { + t.Fatal(err) + } var labels []uint32 for ; lt.IsValid(); lt.Incr() { labels = append(labels, lt.Uint32()) @@ -399,46 +407,7 @@ func TestRelabelPointsInLabelRange(t *testing.T) { } func TestUnlabelPoints(t *testing.T) { - header := pc.PointCloudHeader{ - Fields: []string{"x", "y", "z", "label"}, - Size: []int{4, 4, 4, 4}, - Type: []string{"F", "F", "F", "U"}, - Count: []int{1, 1, 1, 1}, - Width: 4, - Height: 1, - } - pp := &pc.PointCloud{ - PointCloudHeader: header, - Points: 4, - Data: make([]byte, 4*4*4), - } - it, err := pp.Vec3Iterator() - if err != nil { - t.Fatal(err) - } - it.SetVec3(mat.Vec3{1, 2, 3}) - it.Incr() - it.SetVec3(mat.Vec3{4, 5, 6}) - it.Incr() - it.SetVec3(mat.Vec3{7, 8, 9}) - it.Incr() - it.SetVec3(mat.Vec3{10, 11, 12}) - - lt, err := pp.Uint32Iterator("label") - if err != nil { - t.Fatal(err) - } - lt.SetUint32(0) - lt.Incr() - lt.SetUint32(1) - lt.Incr() - lt.SetUint32(2) - lt.Incr() - lt.SetUint32(3) - lt.Incr() - c := newCommandContext(&dummyPCDIO{}, nil) - c.SetPointCloud(pp, cloudMain) testCases := map[string]struct { labelsToKeep []uint32 @@ -461,13 +430,17 @@ func TestUnlabelPoints(t *testing.T) { for name, tt := range testCases { tt := tt t.Run(name, func(t *testing.T) { - c.SetPointCloud(pp, cloudMain) + // Edits mutate the cloud in place; use a fresh one per case. + c.SetPointCloud(newLabeledTestCloud(t), cloudMain) if err := c.UnlabelPoints(tt.labelsToKeep); err != nil { t.Fatal(err) } - lt, err = c.editor.pp.Uint32Iterator("label") + lt, err := c.editor.pp.Uint32Iterator("label") + if err != nil { + t.Fatal(err) + } var labels []uint32 for ; lt.IsValid(); lt.Incr() { labels = append(labels, lt.Uint32()) diff --git a/editor.go b/editor.go index 78bc24a..0c0442f 100644 --- a/editor.go +++ b/editor.go @@ -141,38 +141,50 @@ func (e *editor) SetPointCloud(pp *pc.PointCloud, id cloudID) error { } func (e *editor) label(fn func(int, mat.Vec3) (uint32, bool)) error { - pcNew := &pc.PointCloud{ - PointCloudHeader: e.pp.PointCloudHeader.Clone(), - Points: e.pp.Points, - Data: make([]byte, len(e.pp.Data)), - } - copy(pcNew.Data, e.pp.Data) - - it, err := pcNew.Vec3Iterator() + it, err := e.pp.Vec3Iterator() if err != nil { return err } - itL, err := pcNew.Uint32Iterator("label") + itL, err := e.pp.Uint32Iterator("label") if err != nil { return err } + p := &labelPatch{} i := 0 for it.IsValid() { - l, ok := fn(i, it.Vec3()) - if ok { - itL.SetUint32(l) + if l, ok := fn(i, it.Vec3()); ok { + if old := itL.Uint32(); old != l { + p.indices = append(p.indices, uint32(i)) + p.oldLabels = append(p.oldLabels, old) + itL.SetUint32(l) + } } it.Incr() itL.Incr() i++ } - e.push(&replacePatch{ - header: e.pp.PointCloudHeader.Clone(), - data: e.pp.Data, - }) - e.pp = pcNew - runtime.GC() + e.push(p) + return nil +} + +func (e *editor) mutateLabels(fn func(i int, l uint32) (uint32, bool)) error { + lt, err := e.pp.Uint32Iterator("label") + if err != nil { + return err + } + + p := &labelPatch{} + for i := 0; lt.IsValid(); i++ { + old := lt.Uint32() + if l, ok := fn(i, old); ok && l != old { + p.indices = append(p.indices, uint32(i)) + p.oldLabels = append(p.oldLabels, old) + lt.SetUint32(l) + } + lt.Incr() + } + e.push(p) return nil } @@ -205,85 +217,23 @@ func (e *editor) passThroughByMask(sel []uint32, mask, val uint32) error { } func (e *editor) relabelPointsInLabelRange(minLabel, maxLabel, newLabel uint32) error { - _, err := e.pp.Uint32Iterator("label") - if err != nil { - return err - } - - pcNew := &pc.PointCloud{ - PointCloudHeader: e.pp.PointCloudHeader.Clone(), - Data: make([]byte, len(e.pp.Data)), - Points: e.pp.Points, - } - copy(pcNew.Data, e.pp.Data) - pcNew.Width = e.pp.Width - pcNew.Height = e.pp.Height - - lt, err := pcNew.Uint32Iterator("label") - if err != nil { - return err - } - - for ; lt.IsValid(); lt.Incr() { - l := lt.Uint32() + return e.mutateLabels(func(_ int, l uint32) (uint32, bool) { if l == newLabel || l < minLabel || l > maxLabel { - continue + return 0, false } - lt.SetUint32(newLabel) - } - - e.push(&replacePatch{ - header: e.pp.PointCloudHeader.Clone(), - data: e.pp.Data, + return newLabel, true }) - e.pp = pcNew - runtime.GC() - return nil } func (e *editor) unlabelPoints(labelsToKeep []uint32) error { - _, err := e.pp.Uint32Iterator("label") - if err != nil { - return err - } - - pcNew := &pc.PointCloud{ - PointCloudHeader: e.pp.PointCloudHeader.Clone(), - Data: make([]byte, len(e.pp.Data)), - Points: e.pp.Points, - } - copy(pcNew.Data, e.pp.Data) - pcNew.Width = e.pp.Width - pcNew.Height = e.pp.Height - - lt, err := pcNew.Uint32Iterator("label") - if err != nil { - return err - } - - isInLabelsToKeep := func(l uint32) bool { + return e.mutateLabels(func(_ int, l uint32) (uint32, bool) { for _, kl := range labelsToKeep { if kl == l { - return true + return 0, false } } - return false - } - - for ; lt.IsValid(); lt.Incr() { - if isInLabelsToKeep(lt.Uint32()) { - continue - } - lt.SetUint32(0) - } - - e.push(&replacePatch{ - header: e.pp.PointCloudHeader.Clone(), - data: e.pp.Data, + return 0, true }) - e.pp = pcNew - runtime.GC() - return nil } func passThrough(pp *pc.PointCloud, fn func(int, mat.Vec3) bool) (*pc.PointCloud, error) { diff --git a/patch.go b/patch.go index 72895de..7032a60 100644 --- a/patch.go +++ b/patch.go @@ -27,8 +27,55 @@ const ( var ( errBrokenPatch = errors.New("broken patch data") errUnknownPatchType = errors.New("unknown patch type") + errNoLabelField = errors.New("point cloud has no label field") ) +func fieldByteOffset(h *pc.PointCloudHeader, name string) (int, bool) { + offset := 0 + for i, fn := range h.Fields { + if fn == name { + return offset, true + } + offset += h.Size[i] * h.Count[i] + } + return 0, false +} + +type labelPatch struct { + indices []uint32 + oldLabels []uint32 +} + +func (p *labelPatch) revert(pp *pc.PointCloud) (*pc.PointCloud, error) { + if len(p.indices) != len(p.oldLabels) { + return nil, errBrokenPatch + } + off, ok := fieldByteOffset(&pp.PointCloudHeader, "label") + if !ok { + return nil, errNoLabelField + } + stride := pp.Stride() + for k, idx := range p.indices { + i := int(idx)*stride + off + if i+4 > len(pp.Data) { + return nil, errBrokenPatch + } + binary.LittleEndian.PutUint32(pp.Data[i:], p.oldLabels[k]) + } + return pp, nil +} + +func (p *labelPatch) encodeHead(buf *bytes.Buffer) { + buf.WriteByte(patchTypeLabel) + writeUint32(buf, uint32(len(p.indices))) + writeUint32s(buf, p.indices) + writeUint32s(buf, p.oldLabels) +} + +func (p *labelPatch) payload() []byte { + return nil +} + type replacePatch struct { header pc.PointCloudHeader data []byte @@ -93,6 +140,16 @@ func decodePatch(b []byte) (patch, []byte, error) { typ := b[0] r := reader{b: b[1:]} switch typ { + case patchTypeLabel: + n := int(r.uint32()) + p := &labelPatch{ + indices: r.uint32s(n), + oldLabels: r.uint32s(n), + } + if r.err != nil { + return nil, nil, r.err + } + return p, r.b, nil case patchTypeReplace: p := &replacePatch{} p.header.Version = math.Float32frombits(r.uint32()) @@ -161,6 +218,14 @@ func writeUint32(buf *bytes.Buffer, v uint32) { buf.Write(b[:]) } +func writeUint32s(buf *bytes.Buffer, vs []uint32) { + b := make([]byte, 4*len(vs)) + for i, v := range vs { + binary.LittleEndian.PutUint32(b[i*4:], v) + } + buf.Write(b) +} + func writeString(buf *bytes.Buffer, s string) { writeUint32(buf, uint32(len(s))) buf.WriteString(s) @@ -184,6 +249,22 @@ func (r *reader) uint32() uint32 { return v } +func (r *reader) uint32s(n int) []uint32 { + if r.err != nil { + return nil + } + if n < 0 || n > len(r.b)/4 { + r.err = errBrokenPatch + return nil + } + vs := make([]uint32, n) + for i := range vs { + vs[i] = binary.LittleEndian.Uint32(r.b[i*4:]) + } + r.b = r.b[4*n:] + return vs +} + func (r *reader) bytes(n int) []byte { if r.err != nil { return nil diff --git a/patch_test.go b/patch_test.go index 1d95637..42abd88 100644 --- a/patch_test.go +++ b/patch_test.go @@ -2,6 +2,7 @@ package main import ( "bytes" + "encoding/binary" "math/rand" "reflect" "testing" @@ -52,6 +53,26 @@ func assertCloudEqual(t *testing.T, expected, got *pc.PointCloud) { } } +func TestLabelPatchRevert(t *testing.T) { + orig := makeTestCloud(t, 100, 100, 1) + pp := cloneCloud(orig) + + stride := pp.Stride() + p := &labelPatch{} + for _, i := range []uint32{0, 3, 42, 99} { + off := int(i)*stride + 12 + p.indices = append(p.indices, i) + p.oldLabels = append(p.oldLabels, binary.LittleEndian.Uint32(pp.Data[off:])) + binary.LittleEndian.PutUint32(pp.Data[off:], 12345) + } + + out, err := p.revert(pp) + if err != nil { + t.Fatal(err) + } + assertCloudEqual(t, orig, out) +} + func TestReplacePatchRevert(t *testing.T) { orig := makeTestCloud(t, 100, 10, 10) orig.Viewpoint = []float32{0, 0, 0, 1, 0, 0, 0} @@ -75,6 +96,7 @@ func TestPatchEncodeDecodeRoundTrip(t *testing.T) { orig := makeTestCloud(t, 100, 10, 10) orig.Viewpoint = []float32{1, 2, 3, 1, 0, 0, 0} patches := []patch{ + &labelPatch{indices: []uint32{1, 2, 42}, oldLabels: []uint32{7, 8, 9}}, &replacePatch{header: orig.PointCloudHeader.Clone(), data: orig.Data}, }