-
Notifications
You must be signed in to change notification settings - Fork 110
/
depth_map_pc.go
169 lines (152 loc) · 4.06 KB
/
depth_map_pc.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
// Package depthadapter is a simple package that turns a DepthMap into a point cloud using intrinsic parameters of a camera.
package depthadapter
import (
"image"
"sync"
"github.com/golang/geo/r3"
"github.com/pkg/errors"
"go.viam.com/utils"
"go.viam.com/rdk/pointcloud"
"go.viam.com/rdk/rimage"
"go.viam.com/rdk/rimage/transform"
)
// ToPointCloud returns a lazy read only pointcloud.
func ToPointCloud(dm *rimage.DepthMap, p transform.Projector) pointcloud.PointCloud {
return newDMPointCloudAdapter(dm, p)
}
const numThreadsDmPointCloudAdapter = 8 // TODO This should probably become a parameter at some point
func newDMPointCloudAdapter(dm *rimage.DepthMap, p transform.Projector) *dmPointCloudAdapter {
var wg sync.WaitGroup
wg.Add(2)
var newDm *rimage.DepthMap
utils.PanicCapturingGo(func() {
defer wg.Done()
newDm = dm.Clone()
})
var size int
sizeChan := make(chan int, numThreadsDmPointCloudAdapter)
utils.PanicCapturingGo(func() {
defer wg.Done()
var sizeWg sync.WaitGroup
sizeWg.Add(numThreadsDmPointCloudAdapter)
// Round up to avoid missing points
batchSize := ((dm.Width() * dm.Height()) + numThreadsDmPointCloudAdapter - 1) / numThreadsDmPointCloudAdapter
for loop := 0; loop < numThreadsDmPointCloudAdapter; loop++ {
f := func(loop int) {
defer sizeWg.Done()
sizeBuf := 0
for i := 0; i < batchSize; i++ {
x := loop*batchSize + i
if x >= dm.Width()*dm.Height() {
break
}
depth := dm.GetDepth(x%dm.Width(), x/dm.Width())
if depth == 0 {
continue
}
sizeBuf++
}
sizeChan <- sizeBuf
}
loopCopy := loop
utils.PanicCapturingGo(func() { f(loopCopy) })
}
sizeWg.Wait()
size = 0
for i := 0; i < numThreadsDmPointCloudAdapter; i++ {
size += <-sizeChan
}
})
wg.Wait()
cache := pointcloud.NewWithPrealloc(size)
return &dmPointCloudAdapter{
dm: newDm,
size: size,
p: p,
cache: cache,
cached: false,
cacheLock: &sync.Mutex{},
}
}
type dmPointCloudAdapter struct {
dm *rimage.DepthMap
p transform.Projector
size int
cache pointcloud.PointCloud
cached bool
cacheLock *sync.Mutex
}
func (dm *dmPointCloudAdapter) safeCacheSet(pt r3.Vector, d pointcloud.Data) error {
dm.cacheLock.Lock()
defer dm.cacheLock.Unlock()
return dm.cache.Set(pt, d)
}
func (dm *dmPointCloudAdapter) Size() int {
return dm.size
}
// genCache generates the cache if it is not already generated.
func (dm *dmPointCloudAdapter) genCache() {
if dm.cached {
return
}
var wg sync.WaitGroup
wg.Add(numThreadsDmPointCloudAdapter)
for loop := 0; loop < numThreadsDmPointCloudAdapter; loop++ {
f := func(loop int) {
defer wg.Done()
// dm.Iterate automatically caches results
dm.Iterate(numThreadsDmPointCloudAdapter, loop, func(p r3.Vector, d pointcloud.Data) bool {
return true
})
}
loopCopy := loop
utils.PanicCapturingGo(func() { f(loopCopy) })
}
wg.Wait()
dm.cached = true
}
func (dm *dmPointCloudAdapter) MetaData() pointcloud.MetaData {
dm.genCache()
return dm.cache.MetaData()
}
func (dm *dmPointCloudAdapter) Set(p r3.Vector, d pointcloud.Data) error {
return errors.New("dmPointCloudAdapter Does not support Set")
}
func (dm *dmPointCloudAdapter) At(x, y, z float64) (pointcloud.Data, bool) {
dm.genCache()
return dm.cache.At(x, y, z)
}
func (dm *dmPointCloudAdapter) Iterate(numBatches, myBatch int, fn func(pt r3.Vector, d pointcloud.Data) bool) {
if dm.cached {
dm.cache.Iterate(numBatches, myBatch, fn)
return
}
for y := 0; y < dm.dm.Height(); y++ {
if numBatches > 0 && y%numBatches != myBatch {
continue
}
for x := 0; x < dm.dm.Width(); x++ {
depth := dm.dm.GetDepth(x, y)
if depth == 0 {
continue
}
vec, err := dm.p.ImagePointTo3DPoint(image.Point{x, y}, depth)
if err != nil {
panic(err)
}
if !dm.cached {
err = dm.safeCacheSet(vec, nil)
if err != nil {
panic(err)
}
}
if !fn(vec, nil) {
return
}
}
}
// Since there is no orchestrator for Iterate, we need to check within each process
if dm.size == dm.cache.Size() {
dm.cached = true
}
}