-
Notifications
You must be signed in to change notification settings - Fork 109
/
classifier.go
115 lines (104 loc) · 3.66 KB
/
classifier.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
package transformpipeline
import (
"context"
"image"
"strings"
"github.com/pkg/errors"
"go.opencensus.io/trace"
"go.viam.com/rdk/components/camera"
"go.viam.com/rdk/gostream"
"go.viam.com/rdk/resource"
"go.viam.com/rdk/rimage/transform"
"go.viam.com/rdk/robot"
"go.viam.com/rdk/services/vision"
"go.viam.com/rdk/utils"
"go.viam.com/rdk/vision/classification"
)
// classifierConfig is the attribute struct for classifiers.
type classifierConfig struct {
ClassifierName string `json:"classifier_name"`
ConfidenceThreshold float64 `json:"confidence_threshold"`
MaxClassifications uint32 `json:"max_classifications"`
ValidLabels []string `json:"valid_labels"`
}
// classifierSource takes an image from the camera, and overlays labels from the classifier.
type classifierSource struct {
stream gostream.VideoStream
classifierName string
maxClassifications uint32
labelFilter classification.Postprocessor
confFilter classification.Postprocessor
r robot.Robot
}
func newClassificationsTransform(
ctx context.Context,
source gostream.VideoSource, r robot.Robot, am utils.AttributeMap,
) (gostream.VideoSource, camera.ImageType, error) {
conf, err := resource.TransformAttributeMap[*classifierConfig](am)
if err != nil {
return nil, camera.UnspecifiedStream, err
}
props, err := propsFromVideoSource(ctx, source)
if err != nil {
return nil, camera.UnspecifiedStream, err
}
var cameraModel transform.PinholeCameraModel
cameraModel.PinholeCameraIntrinsics = props.IntrinsicParams
if props.DistortionParams != nil {
cameraModel.Distortion = props.DistortionParams
}
validLabels := make(map[string]interface{})
for _, l := range conf.ValidLabels {
validLabels[strings.ToLower(l)] = struct{}{}
}
labelFilter := classification.NewLabelFilter(validLabels)
confFilter := classification.NewScoreFilter(conf.ConfidenceThreshold)
var maxClassifications uint32 = 1
if conf.MaxClassifications != 0 {
maxClassifications = conf.MaxClassifications
}
classifier := &classifierSource{
gostream.NewEmbeddedVideoStream(source),
conf.ClassifierName,
maxClassifications,
labelFilter,
confFilter,
r,
}
src, err := camera.NewVideoSourceFromReader(ctx, classifier, &cameraModel, camera.ColorStream)
if err != nil {
return nil, camera.UnspecifiedStream, err
}
return src, camera.ColorStream, err
}
// Read returns the image overlaid with at most max_classifications labels from the classification.
// It overlays the labels along with the confidence scores, as long as the scores are above the
// confidence threshold.
func (cs *classifierSource) Read(ctx context.Context) (image.Image, func(), error) {
ctx, span := trace.StartSpan(ctx, "camera::transformpipeline::classifier::Read")
defer span.End()
srv, err := vision.FromRobot(cs.r, cs.classifierName)
if err != nil {
return nil, nil, errors.Wrap(err, "source_classifier can't find vision service")
}
// get image from source camera
img, release, err := cs.stream.Next(ctx)
if err != nil {
return nil, nil, errors.Wrap(err, "could not get next source image")
}
classifications, err := srv.Classifications(ctx, img, int(cs.maxClassifications), map[string]interface{}{})
if err != nil {
return nil, nil, errors.Wrap(err, "could not get classifications")
}
// overlay labels on the source image
classifications = cs.confFilter(classifications)
classifications = cs.labelFilter(classifications)
res, err := classification.Overlay(img, classifications)
if err != nil {
return nil, nil, errors.Wrap(err, "could not overlay labels")
}
return res, release, nil
}
func (cs *classifierSource) Close(ctx context.Context) error {
return cs.stream.Close(ctx)
}