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 //return < 0 if we need more data
  33 static int get_score(int type, int *seq)
  34 {
  35     switch (type) {
  36     case AV1_OBU_SEQUENCE_HEADER:
  37         *seq = 1;
  38         return -1;
  39     case AV1_OBU_FRAME:
  40     case AV1_OBU_FRAME_HEADER:
  41         return *seq ? AVPROBE_SCORE_EXTENSION + 1 : 0;
  42     case AV1_OBU_METADATA:
  43     case AV1_OBU_PADDING:
  44         return -1;
  45     default:
  46         break;
  47     }
  48     return 0;
  49 }
  50
  51 static int read_header(AVFormatContext *s, const AVRational *framerate, AVBSFContext **bsf, void *logctx)
  52 {
  53     const AVBitStreamFilter *filter = av_bsf_get_by_name("av1_frame_merge");
  54     AVStream *st;
  55     int ret;
  56
  57     if (!filter) {
  58         av_log(logctx, AV_LOG_ERROR, "av1_frame_merge bitstream filter "
  59                "not found. This is a bug, please report it.\n");
  60         return AVERROR_BUG;
  61     }
  62
  63     st = avformat_new_stream(s, NULL);
  64     if (!st)
  65         return AVERROR(ENOMEM);
  66
  67     st->codecpar->codec_type = AVMEDIA_TYPE_VIDEO;
  68     st->codecpar->codec_id = AV_CODEC_ID_AV1;
  69     st->need_parsing = AVSTREAM_PARSE_HEADERS;
  70
  71     st->internal->avctx->framerate = *framerate;
  72     // taken from rawvideo demuxers
  73     avpriv_set_pts_info(st, 64, 1, 1200000);
  74
  75     ret = av_bsf_alloc(filter, bsf);
  76     if (ret < 0)
  77         return ret;
  78
  79     ret = avcodec_parameters_copy((*bsf)->par_in, st->codecpar);
  80     if (ret < 0) {
  81         av_bsf_free(bsf);
  82         return ret;
  83     }
  84
  85     ret = av_bsf_init(*bsf);
  86     if (ret < 0)
  87         av_bsf_free(bsf);
  88
  89     return ret;
  90
  91 }
  92
  93 #define DEC AV_OPT_FLAG_DECODING_PARAM
  94
  95 #if CONFIG_AV1_DEMUXER
  96 typedef struct AnnexBContext {
  97     const AVClass *class;
  98     AVBSFContext *bsf;
  99     uint32_t temporal_unit_size;
 100     uint32_t frame_unit_size;
 101     AVRational framerate;
 102 } AnnexBContext;
 103
 104 static int leb(AVIOContext *pb, uint32_t *len) {
 105     int more, i = 0;
 106     uint8_t byte;
 107     *len = 0;
 108     do {
 109         unsigned bits;
 110         byte = avio_r8(pb);
 111         more = byte & 0x80;
 112         bits = byte & 0x7f;
 113         if (i <= 3 || (i == 4 && bits < (1 << 4)))
 114             *len |= bits << (i * 7);
 115         else if (bits)
 116             return AVERROR_INVALIDDATA;
 117         if (++i == 8 && more)
 118             return AVERROR_INVALIDDATA;
 119         if (pb->eof_reached || pb->error)
 120             return pb->error ? pb->error : AVERROR(EIO);
 121     } while (more);
 122     return i;
 123 }
 124
 125 static int read_obu(const uint8_t *buf, int size, int64_t *obu_size, int *type)
 126 {
 127     int start_pos, temporal_id, spatial_id;
 128     int len;
 129
 130     len = parse_obu_header(buf, size, obu_size, &start_pos,
 131                            type, &temporal_id, &spatial_id);
 132     if (len < 0)
 133         return len;
 134
 135     return 0;
 136 }
 137
 138 static int annexb_probe(const AVProbeData *p)
 139 {
 140     AVIOContext pb;
 141     int64_t obu_size;
 142     uint32_t temporal_unit_size, frame_unit_size, obu_unit_size;
 143     int seq = 0;
 144     int ret, type, cnt = 0;
 145
 146     ffio_init_context(&pb, p->buf, p->buf_size, 0,
 147                       NULL, NULL, NULL, NULL);
 148
 149     ret = leb(&pb, &temporal_unit_size);
 150     if (ret < 0)
 151         return 0;
 152     cnt += ret;
 153     ret = leb(&pb, &frame_unit_size);
 154     if (ret < 0 || ((int64_t)frame_unit_size + ret) > temporal_unit_size)
 155         return 0;
 156     cnt += ret;
 157     ret = leb(&pb, &obu_unit_size);
 158     if (ret < 0 || ((int64_t)obu_unit_size + ret) >= frame_unit_size)
 159         return 0;
 160     cnt += ret;
 161
 162     frame_unit_size -= obu_unit_size + ret;
 163
 164     avio_skip(&pb, obu_unit_size);
 165     if (pb.eof_reached || pb.error)
 166         return 0;
 167
 168     // Check that the first OBU is a Temporal Delimiter.
 169     ret = read_obu(p->buf + cnt, FFMIN(p->buf_size - cnt, obu_unit_size), &obu_size, &type);
 170     if (ret < 0 || type != AV1_OBU_TEMPORAL_DELIMITER || obu_size > 0)
 171         return 0;
 172     cnt += obu_unit_size;
 173
 174     do {
 175         ret = leb(&pb, &obu_unit_size);
 176         if (ret < 0 || ((int64_t)obu_unit_size + ret) > frame_unit_size)
 177             return 0;
 178         cnt += ret;
 179
 180         avio_skip(&pb, obu_unit_size);
 181         if (pb.eof_reached || pb.error)
 182             return 0;
 183
 184         ret = read_obu(p->buf + cnt, FFMIN(p->buf_size - cnt, obu_unit_size), &obu_size, &type);
 185         if (ret < 0)
 186             return 0;
 187         cnt += obu_unit_size;
 188
 189         ret = get_score(type, &seq);
 190         if (ret >= 0)
 191             return ret;
 192
 193         frame_unit_size -= obu_unit_size + ret;
 194     } while (frame_unit_size);
 195
 196     return 0;
 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 #define OFFSET(x) offsetof(AnnexBContext, x)
 271 static const AVOption annexb_options[] = {
 272     { "framerate", "", OFFSET(framerate), AV_OPT_TYPE_VIDEO_RATE, {.str = "25"}, 0, INT_MAX, DEC},
 273     { NULL },
 274 };
 275 #undef OFFSET
 276
 277 static const AVClass annexb_demuxer_class = {
 278     .class_name = "AV1 Annex B demuxer",
 279     .item_name  = av_default_item_name,
 280     .option     = annexb_options,
 281     .version    = LIBAVUTIL_VERSION_INT,
 282 };
 283
 284 const AVInputFormat ff_av1_demuxer = {
 285     .name           = "av1",
 286     .long_name      = NULL_IF_CONFIG_SMALL("AV1 Annex B"),
 287     .priv_data_size = sizeof(AnnexBContext),
 288     .read_probe     = annexb_probe,
 289     .read_header    = annexb_read_header,
 290     .read_packet    = annexb_read_packet,
 291     .read_close     = annexb_read_close,
 292     .extensions     = "obu",
 293     .flags          = AVFMT_GENERIC_INDEX,
 294     .priv_class     = &annexb_demuxer_class,
 295 };
 296 #endif
 297
 298 #if CONFIG_OBU_DEMUXER
 299 typedef struct ObuContext {
 300     const AVClass *class;
 301     AVBSFContext *bsf;
 302     AVRational framerate;
 303     AVFifoBuffer *fifo;
 304 } ObuContext;
 305
 306 //For low overhead obu, we can't foresee the obu size before we parsed the header.
 307 //So, we can't use parse_obu_header here, since it will check size <= buf_size
 308 //see c27c7b49dc for more details
 309 static int read_obu_with_size(const uint8_t *buf, int buf_size, int64_t *obu_size, int *type)
 310 {
 311     GetBitContext gb;
 312     int ret, extension_flag, start_pos;
 313     int64_t size;
 314
 315     ret = init_get_bits8(&gb, buf, FFMIN(buf_size, MAX_OBU_HEADER_SIZE));
 316     if (ret < 0)
 317         return ret;
 318
 319     if (get_bits1(&gb) != 0) // obu_forbidden_bit
 320         return AVERROR_INVALIDDATA;
 321
 322     *type      = get_bits(&gb, 4);
 323     extension_flag = get_bits1(&gb);
 324     if (!get_bits1(&gb))    // has_size_flag
 325         return AVERROR_INVALIDDATA;
 326     skip_bits1(&gb);        // obu_reserved_1bit
 327
 328     if (extension_flag) {
 329         get_bits(&gb, 3);   // temporal_id
 330         get_bits(&gb, 2);   // spatial_id
 331         skip_bits(&gb, 3);  // extension_header_reserved_3bits
 332     }
 333
 334     *obu_size  = leb128(&gb);
 335     if (*obu_size > INT_MAX)
 336         return AVERROR_INVALIDDATA;
 337
 338     if (get_bits_left(&gb) < 0)
 339         return AVERROR_INVALIDDATA;
 340
 341     start_pos = get_bits_count(&gb) / 8;
 342
 343     size = *obu_size + start_pos;
 344     if (size > INT_MAX)
 345         return AVERROR_INVALIDDATA;
 346     return size;
 347 }
 348
 349 static int obu_probe(const AVProbeData *p)
 350 {
 351     int64_t obu_size;
 352     int seq = 0;
 353     int ret, type, cnt;
 354
 355     // Check that the first OBU is a Temporal Delimiter.
 356     cnt = read_obu_with_size(p->buf, p->buf_size, &obu_size, &type);
 357     if (cnt < 0 || type != AV1_OBU_TEMPORAL_DELIMITER || obu_size != 0)
 358         return 0;
 359
 360     while (1) {
 361         ret = read_obu_with_size(p->buf + cnt, p->buf_size - cnt, &obu_size, &type);
 362         if (ret < 0 || obu_size <= 0)
 363             return 0;
 364         cnt += FFMIN(ret, p->buf_size - cnt);
 365
 366         ret = get_score(type, &seq);
 367         if (ret >= 0)
 368             return ret;
 369     }
 370     return 0;
 371 }
 372
 373 static int obu_read_header(AVFormatContext *s)
 374 {
 375     ObuContext *c = s->priv_data;
 376     c->fifo = av_fifo_alloc(MAX_OBU_HEADER_SIZE);
 377     if (!c->fifo)
 378         return AVERROR(ENOMEM);
 379     return read_header(s, &c->framerate, &c->bsf, c);
 380 }
 381
 382 static int obu_get_packet(AVFormatContext *s, AVPacket *pkt)
 383 {
 384     ObuContext *c = s->priv_data;
 385     uint8_t header[MAX_OBU_HEADER_SIZE + AV_INPUT_BUFFER_PADDING_SIZE];
 386     int64_t obu_size;
 387     int size = av_fifo_space(c->fifo);
 388     int ret, len, type;
 389
 390     av_fifo_generic_write(c->fifo, s->pb, size,
 391                           (int (*)(void*, void*, int))avio_read);
 392     size = av_fifo_size(c->fifo);
 393     if (!size)
 394         return 0;
 395
 396     av_fifo_generic_peek(c->fifo, header, size, NULL);
 397
 398     len = read_obu_with_size(header, size, &obu_size, &type);
 399     if (len < 0) {
 400         av_log(c, AV_LOG_ERROR, "Failed to read obu\n");
 401         return len;
 402     }
 403
 404     ret = av_new_packet(pkt, len);
 405     if (ret < 0) {
 406         av_log(c, AV_LOG_ERROR, "Failed to allocate packet for obu\n");
 407         return ret;
 408     }
 409     size = FFMIN(size, len);
 410     av_fifo_generic_read(c->fifo, pkt->data, size, NULL);
 411     len -= size;
 412     if (len > 0) {
 413         ret = avio_read(s->pb, pkt->data + size, len);
 414         if (ret != len) {
 415             av_log(c, AV_LOG_ERROR, "Failed to read %d frome file\n", len);
 416             return ret < 0 ? ret : AVERROR_INVALIDDATA;
 417         }
 418     }
 419     return 0;
 420 }
 421
 422 static int obu_read_packet(AVFormatContext *s, AVPacket *pkt)
 423 {
 424     ObuContext *c = s->priv_data;
 425     int ret;
 426
 427     while (1) {
 428         ret = obu_get_packet(s, pkt);
 429         if (ret < 0)
 430             return ret;
 431         ret = av_bsf_send_packet(c->bsf, pkt);
 432         if (ret < 0) {
 433             av_log(s, AV_LOG_ERROR, "Failed to send packet to "
 434                                     "av1_frame_merge filter\n");
 435             return ret;
 436         }
 437         ret = av_bsf_receive_packet(c->bsf, pkt);
 438         if (ret < 0 && ret != AVERROR(EAGAIN) && ret != AVERROR_EOF)
 439             av_log(s, AV_LOG_ERROR, "av1_frame_merge filter failed to "
 440                                     "send output packet\n");
 441         if (ret != AVERROR(EAGAIN))
 442             break;
 443     }
 444
 445     return ret;
 446 }
 447
 448 static int obu_read_close(AVFormatContext *s)
 449 {
 450     ObuContext *c = s->priv_data;
 451
 452     av_fifo_freep(&c->fifo);
 453     av_bsf_free(&c->bsf);
 454     return 0;
 455 }
 456
 457 #define OFFSET(x) offsetof(ObuContext, x)
 458 static const AVOption obu_options[] = {
 459     { "framerate", "", OFFSET(framerate), AV_OPT_TYPE_VIDEO_RATE, {.str = "25"}, 0, INT_MAX, DEC},
 460     { NULL },
 461 };
 462 #undef OFFSET
 463
 464 static const AVClass obu_demuxer_class = {
 465     .class_name = "AV1 low overhead OBU demuxer",
 466     .item_name  = av_default_item_name,
 467     .option     = obu_options,
 468     .version    = LIBAVUTIL_VERSION_INT,
 469 };
 470
 471 const AVInputFormat ff_obu_demuxer = {
 472     .name           = "obu",
 473     .long_name      = NULL_IF_CONFIG_SMALL("AV1 low overhead OBU"),
 474     .priv_data_size = sizeof(ObuContext),
 475     .read_probe     = obu_probe,
 476     .read_header    = obu_read_header,
 477     .read_packet    = obu_read_packet,
 478     .read_close     = obu_read_close,
 479     .extensions     = "obu",
 480     .flags          = AVFMT_GENERIC_INDEX,
 481     .priv_class     = &obu_demuxer_class,
 482 };
 483 #endif