diff mbox

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

Message ID 20191229140827.88073-1-kaveasaf@gmail.com
State Superseded
Headers show

Commit Message

Asaf Kave Dec. 29, 2019, 2:08 p.m. UTC 
---
 libavcodec/hevc_refs.c |  15 ++++
 libavcodec/hevcdec.c   | 173 ++++++++++++++++++++++++++++++++++++++++-
 libavcodec/hevcdec.h   |  13 ++++
 3 files changed, 200 insertions(+), 1 deletion(-)

Comments

Asaf Kave Jan. 2, 2020, 9:59 a.m. UTC | #1 
On Sun, Dec 29, 2019 at 4:08 PM Asaf Kave <kaveasaf@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
Asaf Kave Jan. 6, 2020, 9:01 a.m. UTC | #2 
On Thu, Jan 2, 2020 at 11:59 AM Asaf Kave <kaveasaf@gmail.com> wrote:

>
>
> On Sun, Dec 29, 2019 at 4:08 PM Asaf Kave <kaveasaf@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
Zhong Li Jan. 7, 2020, 2:48 a.m. UTC | #3 
Haven't tested but patch LGTM

Asaf Kave <kaveasaf@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@ffmpeg.org
> https://ffmpeg.org/mailman/listinfo/ffmpeg-devel
>
> To unsubscribe, visit link above, or email
> ffmpeg-devel-request@ffmpeg.org with subject "unsubscribe".
Asaf Kave Jan. 8, 2020, 9:30 a.m. UTC | #4 
On Tue, Jan 7, 2020 at 4:49 AM Zhong Li <lizhong1008@gmail.com> wrote:

> Haven't tested but patch LGTM
>

Thanks Li for the feedback.


>
> Asaf Kave <kaveasaf@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@ffmpeg.org
> > https://ffmpeg.org/mailman/listinfo/ffmpeg-devel
> >
> > To unsubscribe, visit link above, or email
> > ffmpeg-devel-request@ffmpeg.org with subject "unsubscribe".
>
Asaf Kave Jan. 12, 2020, 8:17 a.m. UTC | #5 
On Wed, Jan 8, 2020 at 11:30 AM Asaf Kave <kaveasaf@gmail.com> wrote:

>
>
> On Tue, Jan 7, 2020 at 4:49 AM Zhong Li <lizhong1008@gmail.com> wrote:
>
>> Haven't tested but patch LGTM
>>
>
> Thanks Li for the feedback.
>
>
>>
>> Asaf Kave <kaveasaf@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@ffmpeg.org
>> > https://ffmpeg.org/mailman/listinfo/ffmpeg-devel
>> >
>> > To unsubscribe, visit link above, or email
>> > ffmpeg-devel-request@ffmpeg.org with subject "unsubscribe".
>>
>
Hello ,
Is there anyone else that can push this to the master branch?
Thanks
Andriy Gelman Jan. 12, 2020, 7:02 p.m. UTC | #6 
Hello Asaf, 

If you compile the code, there are many warning about mixed declaration and code. 
I had a quick look code and have comments below:

On Sun, 29. Dec 16:08, Asaf Kave 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 };

do you need this stack variable? 

> +    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;

avoid camelcase

> +    struct MvField current_mv = {{{ 0 }}};

don't think struct keyword is needed 

> +
> +    /* 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;

you can just break here, instead of checking pFrame == NULL at each iteration of
the loop.

> +        }
> +    }
> +

> +    if(pFrame == NULL) {
> +        av_log(s->avctx, AV_LOG_WARNING,
> +            "Not exporting MVs for frame POC %d", s->poc);
> +        return -1;
> +    }

You are not checking the return of export_mvs() when it's called.
Use an approriate error code, maybe AVERROR_INVALIDDATA ? 

> +
> +    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;
> +    }

Alignment for av_log message and 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;
> +        enum PartMode part_mode = current_cu.cu.part_mode;
> +

> +        int cb_size = 1<<log2_cb_size;

Add space between << operator (in other places in code too)

> +

> +        current_mv = tab_mvf[(y0 >> log2_min_pu_size) * min_pu_width +
> +                       (x0 >> log2_min_pu_size)];

alignment

> +
> +        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;

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 +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);
> +            }

brace alignment and check return for export_mvs()

>          *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@ffmpeg.org
> https://ffmpeg.org/mailman/listinfo/ffmpeg-devel
> 
> To unsubscribe, visit link above, or email
> ffmpeg-devel-request@ffmpeg.org with subject "unsubscribe".
Asaf Kave Jan. 14, 2020, 3:38 p.m. UTC | #7 
On Sun, Jan 12, 2020 at 9:02 PM Andriy Gelman <andriy.gelman@gmail.com>
wrote:

> Hello Asaf,
>
> If you compile the code, there are many warning about mixed declaration
> and code.
>

Hi Andriy,
I will take a look again and try to avoid those warnings.
Also i will fix the points you mention, thank you for your time.
Will update patch soon.


I had a quick look code and have comments below:
>
> On Sun, 29. Dec 16:08, Asaf Kave 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 };
>
> do you need this stack variable?
>

Yes, i am assigning it to the ' cuh' that is alias to 's->ref->cuh' , that
holds  all cuh's for the entire frame, for the exporting method.
I can avoid it, but i think is more clear to understand.
What is your opinion?


> > +    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;
>
> avoid camelcase
>

will change.


> > +    struct MvField current_mv = {{{ 0 }}};
>
> don't think struct keyword is needed
>

will remove.


> > +
> > +    /* 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;
>
> you can just break here, instead of checking pFrame == NULL at each
> iteration of
> the loop.
>

OK, I will change this.


> > +        }
> > +    }
> > +
>
> > +    if(pFrame == NULL) {
> > +        av_log(s->avctx, AV_LOG_WARNING,
> > +            "Not exporting MVs for frame POC %d", s->poc);
> > +        return -1;
> > +    }
>
> You are not checking the return of export_mvs() when it's called.
> Use an approriate error code, maybe AVERROR_INVALIDDATA ?
>

good  point, i will return  AVERROR_INVALIDDATA.


>
> > +
> > +    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;
> > +    }
>
> Alignment for av_log message and return AVERROR(ENOMEM)
>

Will change.


>
> > +
> > +    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;
>
> Add space between << operator (in other places in code too)
>

Will change.


>
> > +
>
> > +        current_mv = tab_mvf[(y0 >> log2_min_pu_size) * min_pu_width +
> > +                       (x0 >> log2_min_pu_size)];
>
> alignment
>

Will change


>
> > +
> > +        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;
>
> return AVERROR(ENOMEM)
>

Will change.


>
> > +        }
> > +        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);
> > +            }
>
> brace alignment and check return for export_mvs()
>

Will change , and add appropriate log.


> >          *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@ffmpeg.org
> > https://ffmpeg.org/mailman/listinfo/ffmpeg-devel
> >
> > To unsubscribe, visit link above, or email
> > ffmpeg-devel-request@ffmpeg.org with subject "unsubscribe".
>
> --
> Andriy
>
diff mbox

Patch

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