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 590c9a0..792a384 100644 --- a/editor.go +++ b/editor.go @@ -131,38 +131,49 @@ 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 := &undoDataLabels{} 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++ } - if err := e.push(newUndoDataEntireCloud(e.pp)); err != nil { + return e.push(p) +} + +func (e *editor) mutateLabels(fn func(i int, l uint32) (uint32, bool)) error { + lt, err := e.pp.Uint32Iterator("label") + if err != nil { return err } - e.pp = pcNew - runtime.GC() - return nil + + p := &undoDataLabels{} + 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() + } + return e.push(p) } func (e *editor) passThrough(fn func(int, mat.Vec3) bool) error { @@ -192,83 +203,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) - } - - if err := e.push(newUndoDataEntireCloud(e.pp)); err != nil { - return err - } - e.pp = pcNew - runtime.GC() - return nil + return newLabel, true + }) } 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) - } - - if err := e.push(newUndoDataEntireCloud(e.pp)); err != nil { - return err - } - e.pp = pcNew - runtime.GC() - return nil + return 0, true + }) } func passThrough(pp *pc.PointCloud, fn func(int, mat.Vec3) bool) (*pc.PointCloud, error) { diff --git a/record.go b/record.go index a5303a9..cf5c040 100644 --- a/record.go +++ b/record.go @@ -2,7 +2,9 @@ package main import ( "bytes" + "encoding/binary" "encoding/gob" + "errors" "io" "github.com/seqsense/pcgol/pc" @@ -18,8 +20,14 @@ type undoData interface { func init() { gob.Register(&undoDataEntireCloud{}) + gob.Register(&undoDataLabels{}) } +var ( + errBrokenRecord = errors.New("broken undo record") + errNoLabelField = errors.New("point cloud has no label field") +) + type undoDataEntireCloud struct { Header pc.PointCloudHeader data []byte @@ -48,6 +56,47 @@ func (p *undoDataEntireCloud) setPayload(data []byte) { p.data = data } +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 undoDataLabels struct { + Indices []uint32 + OldLabels []uint32 +} + +func (p *undoDataLabels) restore(pp *pc.PointCloud) (*pc.PointCloud, error) { + if len(p.Indices) != len(p.OldLabels) { + return nil, errBrokenRecord + } + 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, errBrokenRecord + } + binary.LittleEndian.PutUint32(pp.Data[i:], p.OldLabels[k]) + } + return pp, nil +} + +func (p *undoDataLabels) payload() []byte { + return nil +} + +func (p *undoDataLabels) setPayload([]byte) {} + // A record is this encoding followed by the raw payload. func encodeUndoData(w io.Writer, d undoData) error { return gob.NewEncoder(w).Encode(&d) diff --git a/record_test.go b/record_test.go index 656bd68..2f7c1e0 100644 --- a/record_test.go +++ b/record_test.go @@ -2,6 +2,7 @@ package main import ( "bytes" + "encoding/binary" "math/rand" "reflect" "testing" @@ -53,6 +54,26 @@ func assertCloudEqual(t *testing.T, expected, got *pc.PointCloud) { } } +func TestUndoDataLabelsRestore(t *testing.T) { + orig := makeTestCloud(t, 100, 100, 1) + pp := cloneCloud(orig) + + stride := pp.Stride() + p := &undoDataLabels{} + 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.restore(pp) + if err != nil { + t.Fatal(err) + } + assertCloudEqual(t, orig, out) +} + func TestUndoDataEntireCloudRestore(t *testing.T) { orig := makeTestCloud(t, 100, 10, 10) pp := makeTestCloud(t, 5, 5, 1) @@ -71,18 +92,23 @@ func TestUndoDataEntireCloudRestore(t *testing.T) { func TestRecordEncodeDecodeRoundTrip(t *testing.T) { orig := makeTestCloud(t, 100, 10, 10) orig.Viewpoint = []float32{1, 2, 3, 1, 0, 0, 0} - p := newUndoDataEntireCloud(orig) - - var buf bytes.Buffer - if err := encodeUndoData(&buf, p); err != nil { - t.Fatal(err) - } - buf.Write(p.payload()) - decoded, err := decodeRecord(buf.Bytes()) - if err != nil { - t.Fatal(err) - } - if !reflect.DeepEqual(undoData(p), decoded) { - t.Errorf("Expected %+v, got %+v", p, decoded) + for name, d := range map[string]undoData{ + "EntireCloud": newUndoDataEntireCloud(orig), + "Labels": &undoDataLabels{Indices: []uint32{1, 2, 42}, OldLabels: []uint32{7, 8, 9}}, + } { + t.Run(name, func(t *testing.T) { + var buf bytes.Buffer + if err := encodeUndoData(&buf, d); err != nil { + t.Fatal(err) + } + buf.Write(d.payload()) + decoded, err := decodeRecord(buf.Bytes()) + if err != nil { + t.Fatal(err) + } + if !reflect.DeepEqual(d, decoded) { + t.Errorf("Expected %+v, got %+v", d, decoded) + } + }) } }