[FFmpeg-devel,v2] HEVC: Export motion vectors to frame side data.

Message ID	20200115091908.81946-2-kaveasaf@gmail.com
State	New
Headers	show Return-Path: <ffmpeg-devel-bounces@ffmpeg.org> From: Asaf Kave <kaveasaf@gmail.com> To: ffmpeg-devel@ffmpeg.org Date: Wed, 15 Jan 2020 11:19:08 +0200 Message-Id: <20200115091908.81946-2-kaveasaf@gmail.com> Subject: [FFmpeg-devel] [PATCH v2] HEVC: Export motion vectors to frame side data. Precedence: list Reply-To: FFmpeg development discussions and patches <ffmpeg-devel@ffmpeg.org> Cc: Asaf Kave <kaveasaf@gmail.com> MIME-Version: 1.0 Content-Type: text/plain; charset="utf-8" Content-Transfer-Encoding: base64 Errors-To: ffmpeg-devel-bounces@ffmpeg.org Sender: "ffmpeg-devel" <ffmpeg-devel-bounces@ffmpeg.org>
Series	[FFmpeg-devel,v2] HEVC: Export motion vectors to frame side data. \| expand [FFmpeg-devel,v2] HEVC: Export motion vectors to frame side data.

Context	Check	Description
andriy/ffmpeg-patchwork	success	Make fate finished

diff --git a/libavcodec/hevc_refs.c b/libavcodec/hevc_refs.c index 7870a72fd6..20f028fa73 100644 --- a/libavcodec/hevc_refs.c +++ b/libavcodec/hevc_refs.c @@ -42,6 +42,9 @@ void ff_hevc_unref_frame(HEVCContext *s, HEVCFrame *frame, int flags) av_buffer_unref(&frame->tab_mvf_buf); frame->tab_mvf = NULL; + av_buffer_unref(&frame->cuh_buf); + frame->cuh = NULL; + av_buffer_unref(&frame->rpl_buf); av_buffer_unref(&frame->rpl_tab_buf); frame->rpl_tab = NULL; @@ -101,11 +104,17 @@ static HEVCFrame *alloc_frame(HEVCContext *s) goto fail; frame->tab_mvf = (MvField *)frame->tab_mvf_buf->data; + frame->cuh_buf = av_buffer_pool_get(s->cuh_pool); + if (!frame->cuh_buf) + goto fail; + frame->cuh = (CodingUnitHelper *)frame->cuh_buf->data; + frame->rpl_tab_buf = av_buffer_pool_get(s->rpl_tab_pool); if (!frame->rpl_tab_buf) goto fail; frame->rpl_tab = (RefPicListTab **)frame->rpl_tab_buf->data; frame->ctb_count = s->ps.sps->ctb_width * s->ps.sps->ctb_height; + frame->cu_count = 0; for (j = 0; j < frame->ctb_count; j++) frame->rpl_tab[j] = (RefPicListTab *)frame->rpl_buf->data; @@ -161,6 +170,10 @@ int ff_hevc_set_new_ref(HEVCContext *s, AVFrame **frame, int poc) else ref->flags = HEVC_FRAME_FLAG_SHORT_REF; + if (s->avctx->flags2 & AV_CODEC_FLAG2_EXPORT_MVS) { + ref->flags |= HEVC_FRAME_FLAG_MV; + } + ref->poc = poc; ref->sequence = s->seq_decode; ref->frame->crop_left = s->ps.sps->output_window.left_offset; @@ -216,6 +229,8 @@ int ff_hevc_output_frame(HEVCContext *s, AVFrame *out, int flush) if (ret < 0) return ret; + s->output_frame_poc = frame->poc; + av_log(s->avctx, AV_LOG_DEBUG, "Output frame with POC %d.\n", frame->poc); return 1; diff --git a/libavcodec/hevcdec.c b/libavcodec/hevcdec.c index 19b0cd815d..367745c6ee 100644 --- a/libavcodec/hevcdec.c +++ b/libavcodec/hevcdec.c @@ -32,6 +32,7 @@ #include "libavutil/opt.h" #include "libavutil/pixdesc.h" #include "libavutil/stereo3d.h" +#include "libavutil/motion_vector.h" #include "bswapdsp.h" #include "bytestream.h" @@ -80,6 +81,7 @@ static void pic_arrays_free(HEVCContext *s) av_freep(&s->sh.offset); av_buffer_pool_uninit(&s->tab_mvf_pool); + av_buffer_pool_uninit(&s->cuh_pool); av_buffer_pool_uninit(&s->rpl_tab_pool); } @@ -128,9 +130,11 @@ static int pic_arrays_init(HEVCContext *s, const HEVCSPS *sps) s->tab_mvf_pool = av_buffer_pool_init(min_pu_size * sizeof(MvField), av_buffer_allocz); + s->cuh_pool = av_buffer_pool_init(min_pu_size * sizeof(CodingUnitHelper), + av_buffer_allocz); s->rpl_tab_pool = av_buffer_pool_init(ctb_count * sizeof(RefPicListTab), av_buffer_allocz); - if (!s->tab_mvf_pool || !s->rpl_tab_pool) + if (!s->tab_mvf_pool || !s->rpl_tab_pool || !s->cuh_pool) goto fail; return 0; @@ -1806,6 +1810,7 @@ static void hls_prediction_unit(HEVCContext *s, int x0, int y0, int min_pu_width = s->ps.sps->min_pu_width; MvField *tab_mvf = s->ref->tab_mvf; + CodingUnitHelper *cuh = s->ref->cuh; RefPicList *refPicList = s->ref->refPicList; HEVCFrame *ref0 = NULL, *ref1 = NULL; uint8_t *dst0 = POS(0, x0, y0); @@ -1843,6 +1848,10 @@ static void hls_prediction_unit(HEVCContext *s, int x0, int y0, for (i = 0; i < nPbW >> s->ps.sps->log2_min_pu_size; i++) tab_mvf[(y_pu + j) * min_pu_width + x_pu + i] = current_mv; + cuh[s->ref->cu_count].cu = lc->cu; + cuh[s->ref->cu_count].log2_cu_size = log2_cb_size; + s->ref->cu_count++; + if (current_mv.pred_flag & PF_L0) { ref0 = refPicList[0].ref[current_mv.ref_idx[0]]; if (!ref0) @@ -3192,6 +3201,163 @@ static int hevc_decode_extradata(HEVCContext *s, uint8_t *buf, int length, int f return 0; } +static int set_mv(AVMotionVector *mv, int puW, int puH, + int dst_x, int dst_y, + int motion_x, int motion_y, int motion_scale, + int direction) +{ + mv->w = puW; + mv->h = puH; + mv->motion_x = motion_x; + mv->motion_y = motion_y; + mv->motion_scale = motion_scale; + mv->dst_x = dst_x; + mv->dst_y = dst_y; + mv->src_x = dst_x + motion_x / motion_scale; + mv->src_y = dst_y + motion_y / motion_scale; + mv->source = direction ? 1 : -1; + mv->flags = 0; + + return 1; +} + +static int add_mv(HEVCContext *s, AVMotionVector *mvs, MvField current_mv, int x0, int y0, int width, int height) +{ + int sx, sy, count = 0; + const int scale = 4; + + sx = x0 + (width / 2); + sy = y0 + (height / 2); + + if (current_mv.pred_flag & PF_L0) { + count += set_mv(mvs + count, width, height , sx, sy, current_mv.mv[0].x, current_mv.mv[0].y, scale, 0); + } + + if (current_mv.pred_flag & PF_L1) { + count += set_mv(mvs + count, width, height , sx, sy, current_mv.mv[1].x, current_mv.mv[1].y, scale, 1); + } + + return count; +} + +static int export_mvs(HEVCContext *s, AVFrame *out) +{ + int x0, y0, i, log2_cb_size; + int mv_count = 0; + HEVCFrame* pframe = NULL; + MvField *tab_mvf = NULL; + struct MvField current_mv = {{{ 0 }}}; + enum PartMode part_mode; + int cb_size; + int half_cb_size; + AVMotionVector *mvs = NULL; + + const int min_pu_width = s->ps.sps->min_pu_width; + const unsigned log2_min_pu_size = s->ps.sps->log2_min_pu_size; + + /* Find the next output picture\frame tp export it VMs */ + for (i = 0; i < FF_ARRAY_ELEMS(s->DPB); i++) { + HEVCFrame *frame = &s->DPB[i]; + if (frame->flags & (HEVC_FRAME_FLAG_OUTPUT | HEVC_FRAME_FLAG_MV) && + frame->poc == s->output_frame_poc) { + pframe = frame; + break; + } + } + + if(pframe == NULL) { + av_log(s->avctx, AV_LOG_WARNING, "Not exporting MVs for frame POC %d", s->poc); + return AVERROR_INVALIDDATA; + } + + tab_mvf = pframe->tab_mvf; + + /* size is number of [coding units * 2 * 4] where 2 is for directions and 4 is + * for the maximum number of part mode */ + mvs = av_malloc_array(pframe->cu_count, 2 * 4 * sizeof(AVMotionVector)); + if (!mvs) { + av_log(s->avctx, AV_LOG_WARNING, "Failed to allocate Motion Vector array for frame POC %d", s->poc); + ff_hevc_unref_frame(s, pframe, HEVC_FRAME_FLAG_MV); + return AVERROR(ENOMEM); + } + + for (i = 0; i < pframe->cu_count; ++i) { + + const CodingUnitHelper current_cu = pframe->cuh[i]; + /* Export only INTER prediction coding units */ + if(current_cu.cu.pred_mode != MODE_INTER) + continue; + + y0 = current_cu.cu.y; + x0 = current_cu.cu.x; + log2_cb_size = current_cu.log2_cu_size; + part_mode = current_cu.cu.part_mode; + + cb_size = 1 << log2_cb_size; + + current_mv = tab_mvf[(y0 >> log2_min_pu_size) * min_pu_width + (x0 >> log2_min_pu_size)]; + + half_cb_size = 1 << (log2_cb_size - 1); + switch (part_mode) { + case PART_2Nx2N: + mv_count += add_mv(s, mvs + mv_count, current_mv, x0, y0, cb_size, cb_size); + break; + case PART_NxN: + mv_count += add_mv(s, mvs + mv_count, current_mv, x0, y0, cb_size / 2, cb_size / 2); + mv_count += add_mv(s, mvs + mv_count, current_mv, x0 + half_cb_size, y0, cb_size / 2, cb_size / 2); + mv_count += add_mv(s, mvs + mv_count, current_mv, x0, y0 + half_cb_size, cb_size / 2, cb_size / 2); + mv_count += add_mv(s, mvs + mv_count, current_mv, x0 + half_cb_size, y0 + half_cb_size, cb_size / 2, cb_size / 2); + break; + case PART_2NxN: + mv_count += add_mv(s, mvs + mv_count, current_mv, x0, y0, cb_size, cb_size / 2); + mv_count += add_mv(s, mvs + mv_count, current_mv, x0, y0 + half_cb_size, cb_size, cb_size / 2); + break; + case PART_Nx2N: + mv_count += add_mv(s, mvs + mv_count, current_mv, x0, y0, cb_size / 2, cb_size); + mv_count += add_mv(s, mvs + mv_count, current_mv, x0 + half_cb_size, y0, cb_size / 2, cb_size); + break; + case PART_2NxnU: + mv_count += add_mv(s, mvs + mv_count, current_mv, x0, y0, cb_size, cb_size / 4); + mv_count += add_mv(s, mvs + mv_count, current_mv, x0, y0 + cb_size / 4, cb_size, cb_size * 3 / 4); + break; + case PART_2NxnD: + mv_count += add_mv(s, mvs + mv_count, current_mv, x0, y0, cb_size, cb_size * 3 / 4); + mv_count += add_mv(s, mvs + mv_count, current_mv, x0, y0 + cb_size * 3 / 4, cb_size, cb_size / 4); + break; + case PART_nLx2N: + mv_count += add_mv(s, mvs + mv_count, current_mv, x0, y0, cb_size / 4, cb_size); + mv_count += add_mv(s, mvs + mv_count, current_mv, x0 + cb_size / 4, y0, cb_size * 3 / 4, cb_size); + break; + case PART_nRx2N: + mv_count += add_mv(s, mvs + mv_count, current_mv, x0, y0, cb_size * 3 / 4, cb_size); + mv_count += add_mv(s, mvs + mv_count, current_mv, x0 * 3 / 4, y0, cb_size / 4, cb_size); + break; + default: + break; + } + } + + if (mv_count) { + AVFrameSideData *sd; + + av_log(s->avctx, AV_LOG_DEBUG, "Adding %d MVs info to frame with POC %d", mv_count, s->poc); + sd = av_frame_new_side_data(out, AV_FRAME_DATA_MOTION_VECTORS, mv_count * sizeof(AVMotionVector)); + if (!sd) { + av_log(s->avctx, AV_LOG_WARNING, "Failed to create side data MVs info to frame POC %d", s->poc); + av_freep(&mvs); + ff_hevc_unref_frame(s, pframe, HEVC_FRAME_FLAG_MV); + return AVERROR(ENOMEM); + } + memcpy(sd->data, mvs, mv_count * sizeof(AVMotionVector)); + } + + /* Cleanup and release */ + av_freep(&mvs); + ff_hevc_unref_frame(s, pframe, HEVC_FRAME_FLAG_MV); + + return 0; +} + static int hevc_decode_frame(AVCodecContext *avctx, void *data, int *got_output, AVPacket *avpkt) { @@ -3249,6 +3415,11 @@ static int hevc_decode_frame(AVCodecContext *avctx, void *data, int *got_output, if (s->output_frame->buf[0]) { av_frame_move_ref(data, s->output_frame); + if (s->avctx->flags2 & AV_CODEC_FLAG2_EXPORT_MVS) { + if(export_mvs(s, data) < 0 ) { + av_log(avctx, AV_LOG_WARNING, "Failed to export Motion Vectors for frame with POC %d", s->poc); + } + } *got_output = 1; } @@ -3277,8 +3448,14 @@ static int hevc_ref_frame(HEVCContext *s, HEVCFrame *dst, HEVCFrame *src) if (!dst->rpl_buf) goto fail; + dst->cuh_buf = av_buffer_ref(src->cuh_buf); + if (!dst->cuh_buf) + goto fail; + dst->cuh = src->cuh; + dst->poc = src->poc; dst->ctb_count = src->ctb_count; + dst->cu_count = src->cu_count; dst->flags = src->flags; dst->sequence = src->sequence; diff --git a/libavcodec/hevcdec.h b/libavcodec/hevcdec.h index 89e0809850..187c6ec7a2 100644 --- a/libavcodec/hevcdec.h +++ b/libavcodec/hevcdec.h @@ -252,6 +252,12 @@ typedef struct CodingUnit { uint8_t cu_transquant_bypass_flag; } CodingUnit; +typedef struct CodingUnitHelper { + CodingUnit cu; + + uint8_t log2_cu_size; +} CodingUnitHelper; + typedef struct Mv { int16_t x; ///< horizontal component of motion vector int16_t y; ///< vertical component of motion vector @@ -307,20 +313,24 @@ typedef struct DBParams { #define HEVC_FRAME_FLAG_SHORT_REF (1 << 1) #define HEVC_FRAME_FLAG_LONG_REF (1 << 2) #define HEVC_FRAME_FLAG_BUMPING (1 << 3) +#define HEVC_FRAME_FLAG_MV (1 << 4) typedef struct HEVCFrame { AVFrame *frame; ThreadFrame tf; MvField *tab_mvf; + CodingUnitHelper *cuh; RefPicList *refPicList; RefPicListTab **rpl_tab; int ctb_count; int poc; + int cu_count; struct HEVCFrame *collocated_ref; AVBufferRef *tab_mvf_buf; AVBufferRef *rpl_tab_buf; AVBufferRef *rpl_buf; + AVBufferRef *cuh_buf; AVBufferRef *hwaccel_priv_buf; void *hwaccel_picture_private; @@ -405,11 +415,14 @@ typedef struct HEVCContext { uint8_t *sao_pixel_buffer_h[3]; uint8_t *sao_pixel_buffer_v[3]; + int output_frame_poc; + HEVCParamSets ps; HEVCSEI sei; struct AVMD5 *md5_ctx; AVBufferPool *tab_mvf_pool; + AVBufferPool *cuh_pool; AVBufferPool *rpl_tab_pool; ///< candidate references for the current frame

[FFmpeg-devel,v2] HEVC: Export motion vectors to frame side data.

Checks

Commit Message

Comments

Patch