git.sesse.net Git - ffmpeg/blob - libavformat/av1dec.c

   1 /*
   2  * AV1 Annex B demuxer
   3  * Copyright (c) 2019 James Almer <jamrial@gmail.com>
   4  *
   5  * This file is part of FFmpeg.
   6  *
   7  * FFmpeg is free software; you can redistribute it and/or
   8  * modify it under the terms of the GNU Lesser General Public
   9  * License as published by the Free Software Foundation; either
  10  * version 2.1 of the License, or (at your option) any later version.
  11  *
  12  * FFmpeg is distributed in the hope that it will be useful,
  13  * but WITHOUT ANY WARRANTY; without even the implied warranty of
  14  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  15  * Lesser General Public License for more details.
  16  *
  17  * You should have received a copy of the GNU Lesser General Public
  18  * License along with FFmpeg; if not, write to the Free Software
  19  * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
  20  */
  21
  22 #include "config.h"
  23
  24 #include "libavutil/common.h"
  25 #include "libavutil/fifo.h"
  26 #include "libavutil/opt.h"
  27 #include "libavcodec/av1_parse.h"
  28 #include "avformat.h"
  29 #include "avio_internal.h"
  30 #include "internal.h"
  31
  32 typedef struct AnnexBContext {
  33     const AVClass *class;
  34     AVBSFContext *bsf;
  35     uint32_t temporal_unit_size;
  36     uint32_t frame_unit_size;
  37     AVRational framerate;
  38 } AnnexBContext;
  39
  40 static int leb(AVIOContext *pb, uint32_t *len) {
  41     int more, i = 0;
  42     uint8_t byte;
  43     *len = 0;
  44     do {
  45         unsigned bits;
  46         byte = avio_r8(pb);
  47         more = byte & 0x80;
  48         bits = byte & 0x7f;
  49         if (i <= 3 || (i == 4 && bits < (1 << 4)))
  50             *len |= bits << (i * 7);
  51         else if (bits)
  52             return AVERROR_INVALIDDATA;
  53         if (++i == 8 && more)
  54             return AVERROR_INVALIDDATA;
  55         if (pb->eof_reached || pb->error)
  56             return pb->error ? pb->error : AVERROR(EIO);
  57     } while (more);
  58     return i;
  59 }
  60
  61 static int read_obu(const uint8_t *buf, int size, int64_t *obu_size, int *type)
  62 {
  63     int start_pos, temporal_id, spatial_id;
  64     int len;
  65
  66     len = parse_obu_header(buf, size, obu_size, &start_pos,
  67                            type, &temporal_id, &spatial_id);
  68     if (len < 0)
  69         return len;
  70
  71     return 0;
  72 }
  73
  74 //return < 0 if we need more data
  75 static int get_score(int type, int *seq)
  76 {
  77     switch (type) {
  78     case AV1_OBU_SEQUENCE_HEADER:
  79         *seq = 1;
  80         return -1;
  81     case AV1_OBU_FRAME:
  82     case AV1_OBU_FRAME_HEADER:
  83         return *seq ? AVPROBE_SCORE_EXTENSION + 1 : 0;
  84     case AV1_OBU_METADATA:
  85     case AV1_OBU_PADDING:
  86         return -1;
  87     default:
  88         break;
  89     }
  90     return 0;
  91 }
  92
  93 static int annexb_probe(const AVProbeData *p)
  94 {
  95     AVIOContext pb;
  96     int64_t obu_size;
  97     uint32_t temporal_unit_size, frame_unit_size, obu_unit_size;
  98     int seq = 0;
  99     int ret, type, cnt = 0;
 100
 101     ffio_init_context(&pb, p->buf, p->buf_size, 0,
 102                       NULL, NULL, NULL, NULL);
 103
 104     ret = leb(&pb, &temporal_unit_size);
 105     if (ret < 0)
 106         return 0;
 107     cnt += ret;
 108     ret = leb(&pb, &frame_unit_size);
 109     if (ret < 0 || ((int64_t)frame_unit_size + ret) > temporal_unit_size)
 110         return 0;
 111     cnt += ret;
 112     temporal_unit_size -= ret;
 113     ret = leb(&pb, &obu_unit_size);
 114     if (ret < 0 || ((int64_t)obu_unit_size + ret) >= frame_unit_size)
 115         return 0;
 116     cnt += ret;
 117
 118     temporal_unit_size -= obu_unit_size + ret;
 119     frame_unit_size -= obu_unit_size + ret;
 120
 121     avio_skip(&pb, obu_unit_size);
 122     if (pb.eof_reached || pb.error)
 123         return 0;
 124
 125     // Check that the first OBU is a Temporal Delimiter.
 126     ret = read_obu(p->buf + cnt, FFMIN(p->buf_size - cnt, obu_unit_size), &obu_size, &type);
 127     if (ret < 0 || type != AV1_OBU_TEMPORAL_DELIMITER || obu_size > 0)
 128         return 0;
 129     cnt += obu_unit_size;
 130
 131     do {
 132         ret = leb(&pb, &obu_unit_size);
 133         if (ret < 0 || ((int64_t)obu_unit_size + ret) > frame_unit_size)
 134             return 0;
 135         cnt += ret;
 136
 137         avio_skip(&pb, obu_unit_size);
 138         if (pb.eof_reached || pb.error)
 139             return 0;
 140
 141         ret = read_obu(p->buf + cnt, FFMIN(p->buf_size - cnt, obu_unit_size), &obu_size, &type);
 142         if (ret < 0)
 143             return 0;
 144         cnt += obu_unit_size;
 145
 146         ret = get_score(type, &seq);
 147         if (ret >= 0)
 148             return ret;
 149
 150         temporal_unit_size -= obu_unit_size + ret;
 151         frame_unit_size -= obu_unit_size + ret;
 152     } while (frame_unit_size);
 153
 154     return 0;
 155 }
 156
 157 static int read_header(AVFormatContext *s, const AVRational *framerate, AVBSFContext **bsf, void *logctx)
 158 {
 159     const AVBitStreamFilter *filter = av_bsf_get_by_name("av1_frame_merge");
 160     AVStream *st;
 161     int ret;
 162
 163     if (!filter) {
 164         av_log(logctx, AV_LOG_ERROR, "av1_frame_merge bitstream filter "
 165                "not found. This is a bug, please report it.\n");
 166         return AVERROR_BUG;
 167     }
 168
 169     st = avformat_new_stream(s, NULL);
 170     if (!st)
 171         return AVERROR(ENOMEM);
 172
 173     st->codecpar->codec_type = AVMEDIA_TYPE_VIDEO;
 174     st->codecpar->codec_id = AV_CODEC_ID_AV1;
 175     st->need_parsing = AVSTREAM_PARSE_HEADERS;
 176
 177     st->internal->avctx->framerate = *framerate;
 178     // taken from rawvideo demuxers
 179     avpriv_set_pts_info(st, 64, 1, 1200000);
 180
 181     ret = av_bsf_alloc(filter, bsf);
 182     if (ret < 0)
 183         return ret;
 184
 185     ret = avcodec_parameters_copy((*bsf)->par_in, st->codecpar);
 186     if (ret < 0) {
 187         av_bsf_free(bsf);
 188         return ret;
 189     }
 190
 191     ret = av_bsf_init(*bsf);
 192     if (ret < 0)
 193         av_bsf_free(bsf);
 194
 195     return ret;
 196
 197 }
 198
 199 static int annexb_read_header(AVFormatContext *s)
 200 {
 201     AnnexBContext *c = s->priv_data;
 202     return read_header(s, &c->framerate, &c->bsf, c);
 203 }
 204
 205 static int annexb_read_packet(AVFormatContext *s, AVPacket *pkt)
 206 {
 207     AnnexBContext *c = s->priv_data;
 208     uint32_t obu_unit_size;
 209     int ret, len;
 210
 211 retry:
 212     if (avio_feof(s->pb)) {
 213         if (c->temporal_unit_size || c->frame_unit_size)
 214             return AVERROR(EIO);
 215         goto end;
 216     }
 217
 218     if (!c->temporal_unit_size) {
 219         len = leb(s->pb, &c->temporal_unit_size);
 220         if (len < 0) return AVERROR_INVALIDDATA;
 221     }
 222
 223     if (!c->frame_unit_size) {
 224         len = leb(s->pb, &c->frame_unit_size);
 225         if (len < 0 || ((int64_t)c->frame_unit_size + len) > c->temporal_unit_size)
 226             return AVERROR_INVALIDDATA;
 227         c->temporal_unit_size -= len;
 228     }
 229
 230     len = leb(s->pb, &obu_unit_size);
 231     if (len < 0 || ((int64_t)obu_unit_size + len) > c->frame_unit_size)
 232         return AVERROR_INVALIDDATA;
 233
 234     ret = av_get_packet(s->pb, pkt, obu_unit_size);
 235     if (ret < 0)
 236         return ret;
 237     if (ret != obu_unit_size)
 238         return AVERROR(EIO);
 239
 240     c->temporal_unit_size -= obu_unit_size + len;
 241     c->frame_unit_size -= obu_unit_size + len;
 242
 243 end:
 244     ret = av_bsf_send_packet(c->bsf, pkt);
 245     if (ret < 0) {
 246         av_log(s, AV_LOG_ERROR, "Failed to send packet to "
 247                                 "av1_frame_merge filter\n");
 248         return ret;
 249     }
 250
 251     ret = av_bsf_receive_packet(c->bsf, pkt);
 252     if (ret < 0 && ret != AVERROR(EAGAIN) && ret != AVERROR_EOF)
 253         av_log(s, AV_LOG_ERROR, "av1_frame_merge filter failed to "
 254                                 "send output packet\n");
 255
 256     if (ret == AVERROR(EAGAIN))
 257         goto retry;
 258
 259     return ret;
 260 }
 261
 262 static int annexb_read_close(AVFormatContext *s)
 263 {
 264     AnnexBContext *c = s->priv_data;
 265
 266     av_bsf_free(&c->bsf);
 267     return 0;
 268 }
 269
 270 typedef struct ObuContext {
 271     const AVClass *class;
 272     AVBSFContext *bsf;
 273     AVRational framerate;
 274     AVFifoBuffer *fifo;
 275 } ObuContext;
 276
 277 //For low overhead obu, we can't foresee the obu size before we parsed the header.
 278 //So, we can't use parse_obu_header here, since it will check size <= buf_size
 279 //see c27c7b49dc for more details
 280 static int read_obu_with_size(const uint8_t *buf, int buf_size, int64_t *obu_size, int *type)
 281 {
 282     GetBitContext gb;
 283     int ret, extension_flag, start_pos;
 284     int64_t size;
 285
 286     ret = init_get_bits8(&gb, buf, FFMIN(buf_size, MAX_OBU_HEADER_SIZE));
 287     if (ret < 0)
 288         return ret;
 289
 290     if (get_bits1(&gb) != 0) // obu_forbidden_bit
 291         return AVERROR_INVALIDDATA;
 292
 293     *type      = get_bits(&gb, 4);
 294     extension_flag = get_bits1(&gb);
 295     if (!get_bits1(&gb))    // has_size_flag
 296         return AVERROR_INVALIDDATA;
 297     skip_bits1(&gb);        // obu_reserved_1bit
 298
 299     if (extension_flag) {
 300         get_bits(&gb, 3);   // temporal_id
 301         get_bits(&gb, 2);   // spatial_id
 302         skip_bits(&gb, 3);  // extension_header_reserved_3bits
 303     }
 304
 305     *obu_size  = leb128(&gb);
 306     if (*obu_size > INT_MAX)
 307         return AVERROR_INVALIDDATA;
 308
 309     if (get_bits_left(&gb) < 0)
 310         return AVERROR_INVALIDDATA;
 311
 312     start_pos = get_bits_count(&gb) / 8;
 313
 314     size = *obu_size + start_pos;
 315     if (size > INT_MAX)
 316         return AVERROR_INVALIDDATA;
 317     return size;
 318 }
 319
 320 static int obu_probe(const AVProbeData *p)
 321 {
 322     int64_t obu_size;
 323     int seq = 0;
 324     int ret, type, cnt;
 325
 326     // Check that the first OBU is a Temporal Delimiter.
 327     cnt = read_obu_with_size(p->buf, p->buf_size, &obu_size, &type);
 328     if (cnt < 0 || type != AV1_OBU_TEMPORAL_DELIMITER || obu_size != 0)
 329         return 0;
 330
 331     while (1) {
 332         ret = read_obu_with_size(p->buf + cnt, p->buf_size - cnt, &obu_size, &type);
 333         if (ret < 0 || obu_size <= 0)
 334             return 0;
 335         cnt += ret;
 336
 337         ret = get_score(type, &seq);
 338         if (ret >= 0)
 339             return ret;
 340     }
 341     return 0;
 342 }
 343
 344 static int obu_read_header(AVFormatContext *s)
 345 {
 346     ObuContext *c = s->priv_data;
 347     c->fifo = av_fifo_alloc(MAX_OBU_HEADER_SIZE);
 348     if (!c->fifo)
 349         return AVERROR(ENOMEM);
 350     return read_header(s, &c->framerate, &c->bsf, c);
 351 }
 352
 353 static int obu_prefetch(AVFormatContext *s, uint8_t* dest)
 354 {
 355     ObuContext *c = s->priv_data;
 356     int size = av_fifo_space(c->fifo);
 357     av_fifo_generic_write(c->fifo, s->pb, size,
 358                           (int (*)(void*, void*, int))avio_read);
 359     size = av_fifo_size(c->fifo);
 360     if (size > 0) {
 361         av_fifo_generic_peek(c->fifo, dest, size, NULL);
 362     }
 363     return size;
 364 }
 365
 366 static int obu_read_data(AVFormatContext *s, AVPacket *pkt, int len)
 367 {
 368     int size, left;
 369     ObuContext *c = s->priv_data;
 370     int ret = av_new_packet(pkt, len);
 371     if (ret < 0) {
 372         av_log(c, AV_LOG_ERROR, "Failed to allocate packet for obu\n");
 373         return ret;
 374     }
 375     size = FFMIN(av_fifo_size(c->fifo), len);
 376     av_fifo_generic_read(c->fifo, pkt->data, size, NULL);
 377     left = len - size;
 378     if (left > 0) {
 379         ret = avio_read(s->pb, pkt->data + size, left);
 380         if (ret != left) {
 381             av_log(c, AV_LOG_ERROR, "Failed to read %d frome file\n", left);
 382             return ret;
 383         }
 384     }
 385     return 0;
 386 }
 387
 388 static int obu_get_packet(AVFormatContext *s, AVPacket *pkt)
 389 {
 390     ObuContext *c = s->priv_data;
 391     int64_t obu_size;
 392     int ret, type;
 393     uint8_t header[MAX_OBU_HEADER_SIZE];
 394
 395     ret = obu_prefetch(s, header);
 396     if (!ret)
 397         return AVERROR(EOF);
 398
 399     ret = read_obu_with_size(header, ret, &obu_size, &type);
 400     if (ret < 0) {
 401         av_log(c, AV_LOG_ERROR, "Failed to read obu\n");
 402         return ret;
 403     }
 404     return obu_read_data(s, pkt, ret);
 405 }
 406
 407 static int obu_read_packet(AVFormatContext *s, AVPacket *pkt)
 408 {
 409     ObuContext *c = s->priv_data;
 410     int ret;
 411
 412     while (1) {
 413         ret = obu_get_packet(s, pkt);
 414         if (ret < 0)
 415             return ret;
 416         ret = av_bsf_send_packet(c->bsf, pkt);
 417         if (ret < 0) {
 418             av_log(s, AV_LOG_ERROR, "Failed to send packet to "
 419                                     "av1_frame_merge filter\n");
 420             return ret;
 421         }
 422         ret = av_bsf_receive_packet(c->bsf, pkt);
 423         if (ret < 0 && ret != AVERROR(EAGAIN) && ret != AVERROR_EOF)
 424             av_log(s, AV_LOG_ERROR, "av1_frame_merge filter failed to "
 425                                     "send output packet\n");
 426         if (ret != AVERROR(EAGAIN))
 427             break;
 428     }
 429
 430     return ret;
 431 }
 432
 433 static int obu_read_close(AVFormatContext *s)
 434 {
 435     ObuContext *c = s->priv_data;
 436
 437     av_fifo_freep(&c->fifo);
 438     av_bsf_free(&c->bsf);
 439     return 0;
 440 }
 441
 442 #define DEC AV_OPT_FLAG_DECODING_PARAM
 443
 444 #define OFFSET(x) offsetof(AnnexBContext, x)
 445 static const AVOption annexb_options[] = {
 446     { "framerate", "", OFFSET(framerate), AV_OPT_TYPE_VIDEO_RATE, {.str = "25"}, 0, INT_MAX, DEC},
 447     { NULL },
 448 };
 449 #undef OFFSET
 450
 451 #define OFFSET(x) offsetof(ObuContext, x)
 452 static const AVOption obu_options[] = {
 453     { "framerate", "", OFFSET(framerate), AV_OPT_TYPE_VIDEO_RATE, {.str = "25"}, 0, INT_MAX, DEC},
 454     { NULL },
 455 };
 456 #undef OFFSET
 457
 458 static const AVClass annexb_demuxer_class = {
 459     .class_name = "AV1 Annex B demuxer",
 460     .item_name  = av_default_item_name,
 461     .option     = annexb_options,
 462     .version    = LIBAVUTIL_VERSION_INT,
 463 };
 464
 465 AVInputFormat ff_av1_demuxer = {
 466     .name           = "av1",
 467     .long_name      = NULL_IF_CONFIG_SMALL("AV1 Annex B"),
 468     .priv_data_size = sizeof(AnnexBContext),
 469     .read_probe     = annexb_probe,
 470     .read_header    = annexb_read_header,
 471     .read_packet    = annexb_read_packet,
 472     .read_close     = annexb_read_close,
 473     .extensions     = "obu",
 474     .flags          = AVFMT_GENERIC_INDEX,
 475     .priv_class     = &annexb_demuxer_class,
 476 };
 477
 478 static const AVClass obu_demuxer_class = {
 479     .class_name = "AV1 low overhead OBU demuxer",
 480     .item_name  = av_default_item_name,
 481     .option     = obu_options,
 482     .version    = LIBAVUTIL_VERSION_INT,
 483 };
 484
 485 AVInputFormat ff_obu_demuxer = {
 486     .name           = "obu",
 487     .long_name      = NULL_IF_CONFIG_SMALL("AV1 low overhead OBU"),
 488     .priv_data_size = sizeof(ObuContext),
 489     .read_probe     = obu_probe,
 490     .read_header    = obu_read_header,
 491     .read_packet    = obu_read_packet,
 492     .read_close     = obu_read_close,
 493     .extensions     = "obu",
 494     .flags          = AVFMT_GENERIC_INDEX,
 495     .priv_class     = &obu_demuxer_class,
 496 };