-
Notifications
You must be signed in to change notification settings - Fork 110
/
segmenter.go
121 lines (105 loc) · 3.41 KB
/
segmenter.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
package transformpipeline
import (
"context"
"fmt"
"image"
"go.opencensus.io/trace"
"go.viam.com/rdk/components/camera"
"go.viam.com/rdk/gostream"
"go.viam.com/rdk/pointcloud"
"go.viam.com/rdk/resource"
"go.viam.com/rdk/robot"
"go.viam.com/rdk/services/vision"
"go.viam.com/rdk/spatialmath"
"go.viam.com/rdk/utils"
)
// segmenterConfig is the attribute struct for segementers (their name as found in the vision service).
type segmenterConfig struct {
SegmenterName string `json:"segmenter_name"`
}
// segmenterSource takes a pointcloud from the camera and applies a segmenter to it.
type segmenterSource struct {
stream gostream.VideoStream
cameraName string
segmenterName string
r robot.Robot
}
func newSegmentationsTransform(
ctx context.Context,
source gostream.VideoSource,
r robot.Robot,
am utils.AttributeMap,
sourceString string,
) (gostream.VideoSource, camera.ImageType, error) {
conf, err := resource.TransformAttributeMap[*segmenterConfig](am)
if err != nil {
return nil, camera.UnspecifiedStream, err
}
props, err := propsFromVideoSource(ctx, source)
if err != nil {
return nil, camera.UnspecifiedStream, err
}
segmenter := &segmenterSource{
gostream.NewEmbeddedVideoStream(source),
sourceString,
conf.SegmenterName,
r,
}
src, err := camera.NewVideoSourceFromReader(ctx, segmenter, nil, props.ImageType)
if err != nil {
return nil, camera.UnspecifiedStream, err
}
return src, props.ImageType, err
}
// Validate ensures all parts of the config are valid.
func (cfg *segmenterConfig) Validate(path string) ([]string, error) {
var deps []string
if len(cfg.SegmenterName) == 0 {
return nil, resource.NewConfigValidationFieldRequiredError(path, "segmenter_name")
}
return deps, nil
}
// NextPointCloud function calls a segmenter service on the underlying camera and returns a pointcloud.
func (ss *segmenterSource) NextPointCloud(ctx context.Context) (pointcloud.PointCloud, error) {
ctx, span := trace.StartSpan(ctx, "camera::transformpipeline::segmenter::NextPointCloud")
defer span.End()
// get the service
srv, err := vision.FromRobot(ss.r, ss.segmenterName)
if err != nil {
return nil, fmt.Errorf("source_segmenter cant find vision service: %w", err)
}
// apply service
clouds, err := srv.GetObjectPointClouds(ctx, ss.cameraName, map[string]interface{}{})
if err != nil {
return nil, fmt.Errorf("could not get point clouds: %w", err)
}
if clouds == nil {
return pointcloud.New(), nil
}
// merge pointclouds
cloudsWithOffset := make([]pointcloud.CloudAndOffsetFunc, 0, len(clouds))
for _, cloud := range clouds {
cloudCopy := cloud
cloudFunc := func(ctx context.Context) (pointcloud.PointCloud, spatialmath.Pose, error) {
return cloudCopy, nil, nil
}
cloudsWithOffset = append(cloudsWithOffset, cloudFunc)
}
mergedCloud, err := pointcloud.MergePointClouds(context.Background(), cloudsWithOffset, nil)
if err != nil {
return nil, fmt.Errorf("could not merge point clouds: %w", err)
}
return mergedCloud, nil
}
// Read returns the image if the stream is valid, else error.
func (ss *segmenterSource) Read(ctx context.Context) (image.Image, func(), error) {
img, release, err := ss.stream.Next(ctx)
if err != nil {
return nil, nil, fmt.Errorf("could not get next source image: %w", err)
}
return img, release, nil
}
// Close closes the underlying stream.
func (ss *segmenterSource) Close(ctx context.Context) error {
return ss.stream.Close(ctx)
}