[FFmpeg-devel] [PATCH] HEVC: Export motion vectors to frame side data.
Asaf Kave
kaveasaf at gmail.com
Mon Jan 6 11:01:16 EET 2020
On Thu, Jan 2, 2020 at 11:59 AM Asaf Kave <kaveasaf at gmail.com> wrote:
>
>
> On Sun, Dec 29, 2019 at 4:08 PM Asaf Kave <kaveasaf at gmail.com> wrote:
>
>> ---
>> libavcodec/hevc_refs.c | 15 ++++
>> libavcodec/hevcdec.c | 173 ++++++++++++++++++++++++++++++++++++++++-
>> libavcodec/hevcdec.h | 13 ++++
>> 3 files changed, 200 insertions(+), 1 deletion(-)
>>
>> 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..aedc559283 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,9 @@ 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;
>>
>> + struct CodingUnitHelper cuh_ = {lc->cu, log2_cb_size };
>> + cuh[s->ref->cu_count++] = cuh_;
>> +
>> if (current_mv.pred_flag & PF_L0) {
>> ref0 = refPicList[0].ref[current_mv.ref_idx[0]];
>> if (!ref0)
>> @@ -3192,6 +3200,160 @@ 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;
>> + struct MvField current_mv = {{{ 0 }}};
>> +
>> + /* Find the next output picture\frame tp export it VMs */
>> + for (i = 0; i < FF_ARRAY_ELEMS(s->DPB) && pFrame == NULL; 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;
>> + }
>> + }
>> +
>> + if(pFrame == NULL) {
>> + av_log(s->avctx, AV_LOG_WARNING,
>> + "Not exporting MVs for frame POC %d", s->poc);
>> + return -1;
>> + }
>> +
>> + MvField *tab_mvf = pFrame->tab_mvf;
>> +
>> + const int min_pu_width = s->ps.sps->min_pu_width;
>> + const unsigned log2_min_pu_size = s->ps.sps->log2_min_pu_size;
>> +
>> + /* size is number of [coding units * 2 * 4] where 2 is for
>> directions and 4 is
>> + * for the maximum number of part mode */
>> + AVMotionVector *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 -1;
>> + }
>> +
>> + 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;
>> + enum PartMode part_mode = current_cu.cu.part_mode;
>> +
>> + int cb_size = 1<<log2_cb_size;
>> +
>> + current_mv = tab_mvf[(y0 >> log2_min_pu_size) * min_pu_width +
>> + (x0 >> log2_min_pu_size)];
>> +
>> + const int 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 -1;
>> + }
>> + 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 +3411,9 @@ 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) {
>> + export_mvs(s, data);
>> + }
>> *got_output = 1;
>> }
>>
>> @@ -3277,8 +3442,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
>> --
>> 2.17.1
>>
>>
> Ping
>
2 x ping
More information about the ffmpeg-devel
mailing list