-
Notifications
You must be signed in to change notification settings - Fork 110
/
mods.go
138 lines (124 loc) · 4.29 KB
/
mods.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
package transformpipeline
import (
"context"
"image"
"image/color"
"github.com/disintegration/imaging"
"github.com/edaniels/gostream"
"github.com/pkg/errors"
"go.opencensus.io/trace"
"golang.org/x/image/draw"
"go.viam.com/rdk/components/camera"
"go.viam.com/rdk/config"
"go.viam.com/rdk/rimage"
"go.viam.com/rdk/rimage/transform"
rdkutils "go.viam.com/rdk/utils"
)
// rotateSource is the source to be rotated and the kind of image type.
type rotateSource struct {
originalStream gostream.VideoStream
stream camera.ImageType
}
// newRotateTransform creates a new rotation transform.
func newRotateTransform(ctx context.Context, source gostream.VideoSource, stream camera.ImageType,
) (gostream.VideoSource, camera.ImageType, error) {
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
}
reader := &rotateSource{gostream.NewEmbeddedVideoStream(source), stream}
cam, err := camera.NewFromReader(ctx, reader, &cameraModel, stream)
return cam, stream, err
}
// Next rotates the 2D image depending on the stream type.
func (rs *rotateSource) Read(ctx context.Context) (image.Image, func(), error) {
ctx, span := trace.StartSpan(ctx, "camera::transformpipeline::rotate::Read")
defer span.End()
orig, release, err := rs.originalStream.Next(ctx)
if err != nil {
return nil, nil, err
}
switch rs.stream {
case camera.ColorStream, camera.UnspecifiedStream:
return imaging.Rotate(orig, 180, color.Black), release, nil
case camera.DepthStream:
dm, err := rimage.ConvertImageToDepthMap(ctx, orig)
if err != nil {
return nil, nil, err
}
return dm.Rotate(180), release, nil
default:
return nil, nil, camera.NewUnsupportedImageTypeError(rs.stream)
}
}
// Close closes the original stream.
func (rs *rotateSource) Close(ctx context.Context) error {
return rs.originalStream.Close(ctx)
}
// resizeAttrs are the attributes for a resize transform.
type resizeAttrs struct {
Height int `json:"height_px"`
Width int `json:"width_px"`
}
type resizeSource struct {
originalStream gostream.VideoStream
stream camera.ImageType
height int
width int
}
// newResizeTransform creates a new resize transform.
func newResizeTransform(
ctx context.Context, source gostream.VideoSource, stream camera.ImageType, am config.AttributeMap,
) (gostream.VideoSource, camera.ImageType, error) {
conf, err := config.TransformAttributeMapToStruct(&(resizeAttrs{}), am)
if err != nil {
return nil, camera.UnspecifiedStream, err
}
attrs, ok := conf.(*resizeAttrs)
if !ok {
return nil, camera.UnspecifiedStream, rdkutils.NewUnexpectedTypeError(attrs, conf)
}
if attrs.Width == 0 {
return nil, camera.UnspecifiedStream, errors.New("new width for resize transform cannot be 0")
}
if attrs.Height == 0 {
return nil, camera.UnspecifiedStream, errors.New("new height for resize transform cannot be 0")
}
reader := &resizeSource{gostream.NewEmbeddedVideoStream(source), stream, attrs.Height, attrs.Width}
cam, err := camera.NewFromReader(ctx, reader, nil, stream)
return cam, stream, err
}
// Next resizes the 2D image depending on the stream type.
func (rs *resizeSource) Read(ctx context.Context) (image.Image, func(), error) {
ctx, span := trace.StartSpan(ctx, "camera::transformpipeline::resize::Read")
defer span.End()
orig, release, err := rs.originalStream.Next(ctx)
if err != nil {
return nil, nil, err
}
switch rs.stream {
case camera.ColorStream, camera.UnspecifiedStream:
dst := image.NewRGBA(image.Rect(0, 0, rs.width, rs.height))
draw.NearestNeighbor.Scale(dst, dst.Bounds(), orig, orig.Bounds(), draw.Over, nil)
return dst, release, nil
case camera.DepthStream:
dm, err := rimage.ConvertImageToGray16(orig)
if err != nil {
return nil, nil, err
}
dst := image.NewGray16(image.Rect(0, 0, rs.width, rs.height))
draw.NearestNeighbor.Scale(dst, dst.Bounds(), dm, dm.Bounds(), draw.Over, nil)
return dst, release, nil
default:
return nil, nil, camera.NewUnsupportedImageTypeError(rs.stream)
}
}
// Close closes the original stream.
func (rs *resizeSource) Close(ctx context.Context) error {
return rs.originalStream.Close(ctx)
}