-
Notifications
You must be signed in to change notification settings - Fork 6
/
decoder.go
204 lines (173 loc) · 5.06 KB
/
decoder.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
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
package viamrtsp
/*
#cgo pkg-config: libavcodec libavutil libswscale
#include <libavcodec/avcodec.h>
#include <libavutil/imgutils.h>
#include <libavutil/error.h>
#include <libswscale/swscale.h>
#include <stdlib.h>
*/
import "C"
import (
"fmt"
"image"
"unsafe"
"github.com/pkg/errors"
"go.viam.com/rdk/logging"
)
// decoder is a generic FFmpeg decoder.
type decoder struct {
logger logging.Logger
codecCtx *C.AVCodecContext
srcFrame *C.AVFrame
swsCtx *C.struct_SwsContext
dstFrame *C.AVFrame
dstFramePtr []uint8
}
type videoCodec int
const (
// Unknown indicates an error when no available video codecs could be identified
Unknown videoCodec = iota
// Agnostic indicates that a discrete video codec has yet to be identified
Agnostic
// H264 indicates the h264 video codec
H264
// H265 indicates the h265 video codec
H265
// MJPEG indicates the mjpeg video codec
MJPEG
)
func (vc videoCodec) String() string {
switch vc {
case Agnostic:
return "Agnostic"
case H264:
return "H264"
case H265:
return "H265"
case MJPEG:
return "MJPEG"
default:
return "Unknown"
}
}
func frameData(frame *C.AVFrame) **C.uint8_t {
return (**C.uint8_t)(unsafe.Pointer(&frame.data[0]))
}
func frameLineSize(frame *C.AVFrame) *C.int {
return (*C.int)(unsafe.Pointer(&frame.linesize[0]))
}
// avError converts an AV error code to a AV error message string.
func avError(avErr C.int) string {
var errbuf [C.AV_ERROR_MAX_STRING_SIZE]C.char
if C.av_strerror(avErr, &errbuf[0], C.AV_ERROR_MAX_STRING_SIZE) < 0 {
return fmt.Sprintf("Unknown error with code %d", avErr)
}
return C.GoString(&errbuf[0])
}
// SetLibAVLogLevelFatal sets libav errors to fatal log level
// to cut down on log spam
func SetLibAVLogLevelFatal() {
C.av_log_set_level(C.AV_LOG_FATAL)
}
// newDecoder creates a new decoder for the given codec.
func newDecoder(codecID C.enum_AVCodecID, logger logging.Logger) (*decoder, error) {
codec := C.avcodec_find_decoder(codecID)
if codec == nil {
return nil, errors.New("avcodec_find_decoder() failed")
}
codecCtx := C.avcodec_alloc_context3(codec)
if codecCtx == nil {
return nil, errors.New("avcodec_alloc_context3() failed")
}
res := C.avcodec_open2(codecCtx, codec, nil)
if res < 0 {
C.avcodec_close(codecCtx)
return nil, errors.New("avcodec_open2() failed")
}
srcFrame := C.av_frame_alloc()
if srcFrame == nil {
C.avcodec_close(codecCtx)
return nil, errors.New("av_frame_alloc() failed")
}
return &decoder{
logger: logger,
codecCtx: codecCtx,
srcFrame: srcFrame,
}, nil
}
// newH264Decoder creates a new H264 decoder.
func newH264Decoder(logger logging.Logger) (*decoder, error) {
return newDecoder(C.AV_CODEC_ID_H264, logger)
}
// newH265Decoder creates a new H265 decoder.
func newH265Decoder(logger logging.Logger) (*decoder, error) {
return newDecoder(C.AV_CODEC_ID_H265, logger)
}
// close closes the decoder.
func (d *decoder) close() {
if d.dstFrame != nil {
C.av_frame_free(&d.dstFrame)
}
if d.swsCtx != nil {
C.sws_freeContext(d.swsCtx)
}
C.av_frame_free(&d.srcFrame)
C.avcodec_close(d.codecCtx)
}
func (d *decoder) decode(nalu []byte) (image.Image, error) {
nalu = append(H2645StartCode(), nalu...)
// send frame to decoder
var avPacket C.AVPacket
avPacket.data = (*C.uint8_t)(C.CBytes(nalu))
defer C.free(unsafe.Pointer(avPacket.data))
avPacket.size = C.int(len(nalu))
res := C.avcodec_send_packet(d.codecCtx, &avPacket)
if res < 0 {
return nil, nil
}
// receive frame if available
res = C.avcodec_receive_frame(d.codecCtx, d.srcFrame)
if res < 0 {
return nil, nil
}
// if frame size has changed, allocate needed objects
if d.dstFrame == nil || d.dstFrame.width != d.srcFrame.width || d.dstFrame.height != d.srcFrame.height {
if d.dstFrame != nil {
C.av_frame_free(&d.dstFrame)
}
if d.swsCtx != nil {
C.sws_freeContext(d.swsCtx)
}
d.dstFrame = C.av_frame_alloc()
d.dstFrame.format = C.AV_PIX_FMT_RGBA
d.dstFrame.width = d.srcFrame.width
d.dstFrame.height = d.srcFrame.height
d.dstFrame.color_range = C.AVCOL_RANGE_JPEG
res = C.av_frame_get_buffer(d.dstFrame, 1)
if res < 0 {
return nil, errors.New("av_frame_get_buffer() err")
}
d.swsCtx = C.sws_getContext(d.srcFrame.width, d.srcFrame.height, C.AV_PIX_FMT_YUV420P,
d.dstFrame.width, d.dstFrame.height, (int32)(d.dstFrame.format), C.SWS_BILINEAR, nil, nil, nil)
if d.swsCtx == nil {
return nil, errors.New("sws_getContext() err")
}
dstFrameSize := C.av_image_get_buffer_size((int32)(d.dstFrame.format), d.dstFrame.width, d.dstFrame.height, 1)
d.dstFramePtr = (*[1 << 30]uint8)(unsafe.Pointer(d.dstFrame.data[0]))[:dstFrameSize:dstFrameSize]
}
// convert frame from YUV420 to RGB
res = C.sws_scale(d.swsCtx, frameData(d.srcFrame), frameLineSize(d.srcFrame),
0, d.srcFrame.height, frameData(d.dstFrame), frameLineSize(d.dstFrame))
if res < 0 {
return nil, errors.New("sws_scale() err")
}
// embed frame into an image.Image
return &image.RGBA{
Pix: d.dstFramePtr,
Stride: 4 * (int)(d.dstFrame.width),
Rect: image.Rectangle{
Max: image.Point{(int)(d.dstFrame.width), (int)(d.dstFrame.height)},
},
}, nil
}