1 /*
2 * Copyright 2019 Google Inc.
3 *
4 * Use of this source code is governed by a BSD-style license that can be
5 * found in the LICENSE file.
6 */
7
8 #include "experimental/ffmpeg/SkVideoEncoder.h"
9 #include "include/core/SkColorSpace.h"
10 #include "include/core/SkImage.h"
11 #include "include/private/base/SkTDArray.h"
12
13 extern "C" {
14 #include "libswscale/swscale.h"
15 }
16
17 class SkRandomAccessWStream {
18 SkTDArray<char> fStorage;
19 size_t fPos = 0;
20
21 public:
SkRandomAccessWStream()22 SkRandomAccessWStream() {}
23
pos() const24 size_t pos() const { return fPos; }
25
size() const26 size_t size() const { return fStorage.size(); }
27
write(const void * src,size_t bytes)28 void write(const void* src, size_t bytes) {
29 size_t len = fStorage.size();
30 SkASSERT(fPos <= len);
31
32 size_t overwrite = std::min(len - fPos, bytes);
33 if (overwrite) {
34 SkDebugf("overwrite %zu bytes at %zu offset with %zu remaining\n", overwrite, fPos, bytes - overwrite);
35 memcpy(&fStorage[fPos], src, overwrite);
36 fPos += overwrite;
37 src = (const char*)src + overwrite;
38 bytes -= overwrite;
39 }
40 // bytes now represents the amount to append
41 if (bytes) {
42 fStorage.append(bytes, (const char*)src);
43 fPos += bytes;
44 }
45 SkASSERT(fPos <= fStorage.size());
46 }
47
seek(size_t pos)48 void seek(size_t pos) {
49 SkASSERT(pos <= fStorage.size());
50 fPos = pos;
51 }
52
detachAsData()53 sk_sp<SkData> detachAsData() {
54 // TODO: could add an efficient detach to SkTDArray if we wanted, w/o copy
55 return SkData::MakeWithCopy(fStorage.begin(), fStorage.size());
56 }
57 };
58
59 ///////////////////////////////////////////////////////////////////////////////////////////////////
60
61 // returns true on error (and may dump the particular error message)
check_err(int err,const int silentList[]=nullptr)62 static bool check_err(int err, const int silentList[] = nullptr) {
63 if (err >= 0) {
64 return false;
65 }
66
67 if (silentList) {
68 for (; *silentList; ++silentList) {
69 if (*silentList == err) {
70 return true; // we still report the error, but we don't printf
71 }
72 }
73 }
74
75 char errbuf[128];
76 const char *errbuf_ptr = errbuf;
77
78 if (av_strerror(err, errbuf, sizeof(errbuf)) < 0) {
79 errbuf_ptr = strerror(AVUNERROR(err));
80 }
81 SkDebugf("%s\n", errbuf_ptr);
82 return true;
83 }
84
sk_write_packet(void * ctx,uint8_t * buffer,int size)85 static int sk_write_packet(void* ctx, uint8_t* buffer, int size) {
86 SkRandomAccessWStream* stream = (SkRandomAccessWStream*)ctx;
87 stream->write(buffer, size);
88 return size;
89 }
90
sk_seek_packet(void * ctx,int64_t pos,int whence)91 static int64_t sk_seek_packet(void* ctx, int64_t pos, int whence) {
92 SkRandomAccessWStream* stream = (SkRandomAccessWStream*)ctx;
93 switch (whence) {
94 case SEEK_SET:
95 break;
96 case SEEK_CUR:
97 pos = (int64_t)stream->pos() + pos;
98 break;
99 case SEEK_END:
100 pos = (int64_t)stream->size() + pos;
101 break;
102 default:
103 return -1;
104 }
105 if (pos < 0 || pos > (int64_t)stream->size()) {
106 return -1;
107 }
108 stream->seek(SkToSizeT(pos));
109 return pos;
110 }
111
SkVideoEncoder()112 SkVideoEncoder::SkVideoEncoder() {
113 fInfo = SkImageInfo::MakeUnknown();
114 }
115
~SkVideoEncoder()116 SkVideoEncoder::~SkVideoEncoder() {
117 this->reset();
118
119 if (fSWScaleCtx) {
120 sws_freeContext(fSWScaleCtx);
121 }
122 }
123
reset()124 void SkVideoEncoder::reset() {
125 if (fFrame) {
126 av_frame_free(&fFrame);
127 fFrame = nullptr;
128 }
129 if (fEncoderCtx) {
130 avcodec_free_context(&fEncoderCtx);
131 fEncoderCtx = nullptr;
132 }
133 if (fFormatCtx) {
134 avformat_free_context(fFormatCtx);
135 fFormatCtx = nullptr;
136 }
137
138 av_packet_free(&fPacket);
139 fPacket = nullptr;
140
141 fSurface.reset();
142 fWStream.reset();
143 }
144
init(int fps)145 bool SkVideoEncoder::init(int fps) {
146 // only support this for now
147 AVPixelFormat pix_fmt = AV_PIX_FMT_YUV420P;
148
149 this->reset();
150
151 fWStream.reset(new SkRandomAccessWStream);
152
153 int bufferSize = 4 * 1024;
154 uint8_t* buffer = (uint8_t*)av_malloc(bufferSize);
155 if (!buffer) {
156 return false;
157 }
158 fStreamCtx = avio_alloc_context(buffer, bufferSize, AVIO_FLAG_WRITE, fWStream.get(),
159 nullptr, sk_write_packet, sk_seek_packet);
160 SkASSERT(fStreamCtx);
161
162 avformat_alloc_output_context2(&fFormatCtx, nullptr, "mp4", nullptr);
163 SkASSERT(fFormatCtx);
164 fFormatCtx->pb = fStreamCtx;
165
166 const auto* output_format = fFormatCtx->oformat;
167
168 if (output_format->video_codec == AV_CODEC_ID_NONE) {
169 return false;
170 }
171 const auto* codec = avcodec_find_encoder(output_format->video_codec);
172 SkASSERT(codec);
173
174 fStream = avformat_new_stream(fFormatCtx, codec);
175 SkASSERT(fStream);
176 fStream->id = fFormatCtx->nb_streams-1;
177 fStream->time_base = (AVRational){ 1, fps };
178
179 fEncoderCtx = avcodec_alloc_context3(codec);
180 SkASSERT(fEncoderCtx);
181
182 fEncoderCtx->codec_id = output_format->video_codec;
183 fEncoderCtx->width = fInfo.width();
184 fEncoderCtx->height = fInfo.height();
185 fEncoderCtx->time_base = fStream->time_base;
186 fEncoderCtx->pix_fmt = pix_fmt;
187
188 /* Some formats want stream headers to be separate. */
189 if (output_format->flags & AVFMT_GLOBALHEADER) {
190 fEncoderCtx->flags |= AV_CODEC_FLAG_GLOBAL_HEADER;
191 }
192
193 if (check_err(avcodec_open2(fEncoderCtx, codec, nullptr))) {
194 return false;
195 }
196 fFrame = av_frame_alloc();
197 SkASSERT(fFrame);
198 fFrame->format = pix_fmt;
199 fFrame->width = fEncoderCtx->width;
200 fFrame->height = fEncoderCtx->height;
201 if (check_err(av_frame_get_buffer(fFrame, 32))) {
202 return false;
203 }
204
205 if (check_err(avcodec_parameters_from_context(fStream->codecpar, fEncoderCtx))) {
206 return false;
207 }
208 if (check_err(avformat_write_header(fFormatCtx, nullptr))) {
209 return false;
210 }
211 fPacket = av_packet_alloc();
212 return true;
213 }
214
215 #include "include/core/SkCanvas.h"
216 #include "include/core/SkColorFilter.h"
217 #include "include/core/SkSurface.h"
218 #include "src/core/SkYUVMath.h"
219
is_valid(SkISize dim)220 static bool is_valid(SkISize dim) {
221 if (dim.width() <= 0 || dim.height() <= 0) {
222 return false;
223 }
224 // need the dimensions to be even for YUV 420
225 return ((dim.width() | dim.height()) & 1) == 0;
226 }
227
beginRecording(SkISize dim,int fps)228 bool SkVideoEncoder::beginRecording(SkISize dim, int fps) {
229 if (!is_valid(dim)) {
230 return false;
231 }
232
233 SkAlphaType alphaType = kOpaque_SkAlphaType;
234 sk_sp<SkColorSpace> cs = nullptr; // should we use this?
235 fInfo = SkImageInfo::MakeN32(dim.width(), dim.height(), alphaType, cs);
236 if (!this->init(fps)) {
237 return false;
238 }
239
240 fCurrentPTS = 0;
241 fDeltaPTS = 1;
242
243 const auto fmt = kN32_SkColorType == kRGBA_8888_SkColorType ? AV_PIX_FMT_RGBA : AV_PIX_FMT_BGRA;
244 SkASSERT(sws_isSupportedInput(fmt) > 0);
245 SkASSERT(sws_isSupportedOutput(AV_PIX_FMT_YUV420P) > 0);
246 // sws_getCachedContext takes in either null or a previous ctx. It returns either a new ctx,
247 // or the same as the input if it is compatible with the inputs. Thus we never have to
248 // explicitly release our ctx until the destructor, since sws_getCachedContext takes care
249 // of freeing the old as needed if/when it returns a new one.
250 fSWScaleCtx = sws_getCachedContext(fSWScaleCtx,
251 dim.width(), dim.height(), fmt,
252 dim.width(), dim.height(), AV_PIX_FMT_YUV420P,
253 SWS_FAST_BILINEAR, nullptr, nullptr, nullptr);
254 return fSWScaleCtx != nullptr;
255 }
256
addFrame(const SkPixmap & pm)257 bool SkVideoEncoder::addFrame(const SkPixmap& pm) {
258 if (!is_valid(pm.dimensions())) {
259 return false;
260 }
261 if (pm.info().colorType() != fInfo.colorType()) {
262 return false;
263 }
264 /* make sure the frame data is writable */
265 if (check_err(av_frame_make_writable(fFrame))) {
266 return false;
267 }
268
269 fFrame->pts = fCurrentPTS;
270 fCurrentPTS += fDeltaPTS;
271
272 const uint8_t* src[] = { (const uint8_t*)pm.addr() };
273 const int strides[] = { SkToInt(pm.rowBytes()) };
274 sws_scale(fSWScaleCtx, src, strides, 0, fInfo.height(), fFrame->data, fFrame->linesize);
275
276 return this->sendFrame(fFrame);
277 }
278
sendFrame(AVFrame * frame)279 bool SkVideoEncoder::sendFrame(AVFrame* frame) {
280 if (check_err(avcodec_send_frame(fEncoderCtx, frame))) {
281 return false;
282 }
283
284 int ret = 0;
285 while (ret >= 0) {
286 ret = avcodec_receive_packet(fEncoderCtx, fPacket);
287 if (ret == AVERROR(EAGAIN) || ret == AVERROR_EOF) {
288 break;
289 }
290 if (check_err(ret)) {
291 return false;
292 }
293
294 av_packet_rescale_ts(fPacket, fEncoderCtx->time_base, fStream->time_base);
295 SkASSERT(fPacket->stream_index == fStream->index);
296
297 if (check_err(av_interleaved_write_frame(fFormatCtx, fPacket))) {
298 return false;
299 }
300 }
301 return true;
302 }
303
beginFrame()304 SkCanvas* SkVideoEncoder::beginFrame() {
305 if (!fSurface) {
306 fSurface = SkSurface::MakeRaster(fInfo);
307 if (!fSurface) {
308 return nullptr;
309 }
310 }
311 SkCanvas* canvas = fSurface->getCanvas();
312 canvas->restoreToCount(1);
313 canvas->clear(0);
314 return canvas;
315 }
316
endFrame()317 bool SkVideoEncoder::endFrame() {
318 if (!fSurface) {
319 return false;
320 }
321 SkPixmap pm;
322 return fSurface->peekPixels(&pm) && this->addFrame(pm);
323 }
324
endRecording()325 sk_sp<SkData> SkVideoEncoder::endRecording() {
326 if (!fFormatCtx) {
327 return nullptr;
328 }
329
330 this->sendFrame(nullptr);
331 av_write_trailer(fFormatCtx);
332
333 sk_sp<SkData> data = fWStream->detachAsData();
334 this->reset();
335 return data;
336 }
337