submodule/skia/experimental/ffmpeg/SkVideoEncoder.cpp

   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/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:
  22     SkRandomAccessWStream() {}
  23
  24     size_t pos() const { return fPos; }
  25
  26     size_t size() const { return fStorage.size(); }
  27
  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
  48     void seek(size_t pos) {
  49         SkASSERT(pos <= fStorage.size());
  50         fPos = pos;
  51     }
  52
  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)
  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
  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
  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
 112 SkVideoEncoder::SkVideoEncoder() {
 113     fInfo = SkImageInfo::MakeUnknown();
 114 }
 115
 116 SkVideoEncoder::~SkVideoEncoder() {
 117     this->reset();
 118
 119     if (fSWScaleCtx) {
 120         sws_freeContext(fSWScaleCtx);
 121     }
 122 }
 123
 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
 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
 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
 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
 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
 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
 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
 317 bool SkVideoEncoder::endFrame() {
 318     if (!fSurface) {
 319         return false;
 320     }
 321     SkPixmap pm;
 322     return fSurface->peekPixels(&pm) && this->addFrame(pm);
 323 }
 324
 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 }