From dccd1d2a581606f858e553b9f2a4e8aba1e460b7 Mon Sep 17 00:00:00 2001 From: Andreas Rheinhardt Date: Tue, 3 Aug 2021 00:57:13 +0200 Subject: [PATCH] avformat/yuv4mpegenc: Write data generically via AVPixFmtDescriptor Signed-off-by: Andreas Rheinhardt --- libavformat/yuv4mpegenc.c | 87 +++++++-------------------------------- 1 file changed, 14 insertions(+), 73 deletions(-) diff --git a/libavformat/yuv4mpegenc.c b/libavformat/yuv4mpegenc.c index efa05133d5..fca0ee3120 100644 --- a/libavformat/yuv4mpegenc.c +++ b/libavformat/yuv4mpegenc.c @@ -181,9 +181,8 @@ static int yuv4_write_packet(AVFormatContext *s, AVPacket *pkt) AVStream *st = s->streams[pkt->stream_index]; AVIOContext *pb = s->pb; const AVFrame *frame = (const AVFrame *)pkt->data; - int width, height, h_chroma_shift, v_chroma_shift; - int i; - const uint8_t *ptr, *ptr1, *ptr2; + int width, height; + const AVPixFmtDescriptor *desc; /* construct frame header */ @@ -191,79 +190,21 @@ static int yuv4_write_packet(AVFormatContext *s, AVPacket *pkt) width = st->codecpar->width; height = st->codecpar->height; + desc = av_pix_fmt_desc_get(st->codecpar->format); - ptr = frame->data[0]; + /* The following code presumes all planes to be non-interleaved. */ + for (int k = 0; k < desc->nb_components; k++) { + int plane_height = height, plane_width = width * desc->comp[k].step; + const uint8_t *ptr = frame->data[k]; - switch (st->codecpar->format) { - case AV_PIX_FMT_GRAY8: - case AV_PIX_FMT_YUV411P: - case AV_PIX_FMT_YUV420P: - case AV_PIX_FMT_YUV422P: - case AV_PIX_FMT_YUV444P: - case AV_PIX_FMT_YUVA444P: - // TODO: remove YUVJ pixel formats when they are completely removed from the codebase. - case AV_PIX_FMT_YUVJ420P: - case AV_PIX_FMT_YUVJ422P: - case AV_PIX_FMT_YUVJ444P: - break; - case AV_PIX_FMT_GRAY9: - case AV_PIX_FMT_GRAY10: - case AV_PIX_FMT_GRAY12: - case AV_PIX_FMT_GRAY16: - case AV_PIX_FMT_YUV420P9: - case AV_PIX_FMT_YUV422P9: - case AV_PIX_FMT_YUV444P9: - case AV_PIX_FMT_YUV420P10: - case AV_PIX_FMT_YUV422P10: - case AV_PIX_FMT_YUV444P10: - case AV_PIX_FMT_YUV420P12: - case AV_PIX_FMT_YUV422P12: - case AV_PIX_FMT_YUV444P12: - case AV_PIX_FMT_YUV420P14: - case AV_PIX_FMT_YUV422P14: - case AV_PIX_FMT_YUV444P14: - case AV_PIX_FMT_YUV420P16: - case AV_PIX_FMT_YUV422P16: - case AV_PIX_FMT_YUV444P16: - width *= 2; - break; - default: - av_log(s, AV_LOG_ERROR, "The pixel format '%s' is not supported.\n", - av_get_pix_fmt_name(st->codecpar->format)); - return AVERROR(EINVAL); - } - - for (i = 0; i < height; i++) { - avio_write(pb, ptr, width); - ptr += frame->linesize[0]; - } - - if (st->codecpar->format != AV_PIX_FMT_GRAY8 && st->codecpar->format != AV_PIX_FMT_GRAY9 && - st->codecpar->format != AV_PIX_FMT_GRAY10 && st->codecpar->format != AV_PIX_FMT_GRAY12 && - st->codecpar->format != AV_PIX_FMT_GRAY16) { - // Adjust for smaller Cb and Cr planes - av_pix_fmt_get_chroma_sub_sample(st->codecpar->format, &h_chroma_shift, - &v_chroma_shift); - // Shift right, rounding up - width = AV_CEIL_RSHIFT(width, h_chroma_shift); - height = AV_CEIL_RSHIFT(height, v_chroma_shift); - - ptr1 = frame->data[1]; - ptr2 = frame->data[2]; - for (i = 0; i < height; i++) { /* Cb */ - avio_write(pb, ptr1, width); - ptr1 += frame->linesize[1]; + if (desc->nb_components >= 3 && (k == 1 || k == 2)) { /* chroma? */ + plane_width = AV_CEIL_RSHIFT(plane_width, desc->log2_chroma_w); + plane_height = AV_CEIL_RSHIFT(plane_height, desc->log2_chroma_h); } - for (i = 0; i < height; i++) { /* Cr */ - avio_write(pb, ptr2, width); - ptr2 += frame->linesize[2]; - } - if (st->codecpar->format == AV_PIX_FMT_YUVA444P) { - ptr = frame->data[3]; - for (i = 0; i < height; i++) { /* A */ - avio_write(pb, ptr, width); - ptr += frame->linesize[3]; - } + + for (int i = 0; i < plane_height; i++) { + avio_write(pb, ptr, plane_width); + ptr += frame->linesize[k]; } }