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

Zhong Li lizhong1008 at gmail.com
Tue Jan 7 04:48:49 EET 2020


Haven't tested but patch LGTM

Asaf Kave <kaveasaf at gmail.com> 于2019年12月29日周日 下午10:08写道:
>
> ---
>  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
>
> _______________________________________________
> ffmpeg-devel mailing list
> ffmpeg-devel at ffmpeg.org
> https://ffmpeg.org/mailman/listinfo/ffmpeg-devel
>
> To unsubscribe, visit link above, or email
> ffmpeg-devel-request at ffmpeg.org with subject "unsubscribe".


More information about the ffmpeg-devel mailing list