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