diff mbox

[FFmpeg-devel] libavcodec/libaomenc.c: Added code for computing PSNR/SSIM for libaom encoder.Updated the patch to read the AV_PICTURE_TYPE_I flag for AOM.

Message ID 20180928220905.155627-1-samjohn@google.com
State New
Headers show

Commit Message

Sam John Sept. 28, 2018, 10:09 p.m. UTC
---
 libavcodec/libaomenc.c | 60 +++++++++++++++++++++++++++++++++++++-----
 1 file changed, 54 insertions(+), 6 deletions(-)

Comments

Sam John Oct. 12, 2018, 12:49 a.m. UTC | #1
James,

Is there anything to be done for this patch? Can you please approve this
change and apply the patch to ffmpeg head.



On Fri, Sep 28, 2018 at 3:09 PM Sam John <samjohn@google.com> wrote:

> ---
>  libavcodec/libaomenc.c | 60 +++++++++++++++++++++++++++++++++++++-----
>  1 file changed, 54 insertions(+), 6 deletions(-)
>
> diff --git a/libavcodec/libaomenc.c b/libavcodec/libaomenc.c
> index 6a79d9b873..481338406c 100644
> --- a/libavcodec/libaomenc.c
> +++ b/libavcodec/libaomenc.c
> @@ -50,6 +50,9 @@ struct FrameListData {
>      unsigned long duration;          /**< duration to show frame
>                                            (in timebase units) */
>      uint32_t flags;                  /**< flags for this frame */
> +    uint64_t sse[4];
> +    int have_sse;                    /**< true if we have pending sse[] */
> +    uint64_t frame_number;
>      struct FrameListData *next;
>  };
>
> @@ -68,6 +71,9 @@ typedef struct AOMEncoderContext {
>      int static_thresh;
>      int drop_threshold;
>      int noise_sensitivity;
> +    uint64_t sse[4];
> +    int have_sse; /**< true if we have pending sse[] */
> +    uint64_t frame_number;
>  } AOMContext;
>
>  static const char *const ctlidstr[] = {
> @@ -289,7 +295,8 @@ static av_cold int aom_init(AVCodecContext *avctx,
>  {
>      AOMContext *ctx = avctx->priv_data;
>      struct aom_codec_enc_cfg enccfg = { 0 };
> -    aom_codec_flags_t flags = 0;
> +    aom_codec_flags_t flags =
> +        (avctx->flags & AV_CODEC_FLAG_PSNR) ? AOM_CODEC_USE_PSNR : 0;
>      AVCPBProperties *cpb_props;
>      int res;
>      aom_img_fmt_t img_fmt;
> @@ -499,13 +506,23 @@ static av_cold int aom_init(AVCodecContext *avctx,
>  }
>
>  static inline void cx_pktcpy(struct FrameListData *dst,
> -                             const struct aom_codec_cx_pkt *src)
> +                             const struct aom_codec_cx_pkt *src,
> +                             AOMContext *ctx)
>  {
>      dst->pts      = src->data.frame.pts;
>      dst->duration = src->data.frame.duration;
>      dst->flags    = src->data.frame.flags;
>      dst->sz       = src->data.frame.sz;
>      dst->buf      = src->data.frame.buf;
> +    dst->have_sse = 0;
> +    dst->frame_number = ++ctx->frame_number;
> +    dst->have_sse = ctx->have_sse;
> +    if (ctx->have_sse) {
> +        /* associate last-seen SSE to the frame. */
> +        /* Transfers ownership from ctx to dst. */
> +        memcpy(dst->sse, ctx->sse, sizeof(dst->sse));
> +        ctx->have_sse = 0;
> +    }
>  }
>
>  /**
> @@ -519,6 +536,7 @@ static int storeframe(AVCodecContext *avctx, struct
> FrameListData *cx_frame,
>                        AVPacket *pkt)
>  {
>      AOMContext *ctx = avctx->priv_data;
> +    int pict_type;
>      int ret = ff_alloc_packet2(avctx, pkt, cx_frame->sz, 0);
>      if (ret < 0) {
>          av_log(avctx, AV_LOG_ERROR,
> @@ -528,8 +546,29 @@ static int storeframe(AVCodecContext *avctx, struct
> FrameListData *cx_frame,
>      memcpy(pkt->data, cx_frame->buf, pkt->size);
>      pkt->pts = pkt->dts = cx_frame->pts;
>
> -    if (!!(cx_frame->flags & AOM_FRAME_IS_KEY))
> +    if (!!(cx_frame->flags & AOM_FRAME_IS_KEY)) {
> +        pict_type = AV_PICTURE_TYPE_I;
>          pkt->flags |= AV_PKT_FLAG_KEY;
> +    } else {
> +#ifdef AOM_FRAME_IS_INTRAONLY
> +        if (cx_frame->flags & AOM_FRAME_IS_INTRAONLY)
> +            pict_type = AV_PICTURE_TYPE_I;
> +        else
> +#else
> +            pict_type = AV_PICTURE_TYPE_P;
> +#endif
> +    }
> +
> +    ff_side_data_set_encoder_stats(pkt, 0, cx_frame->sse + 1,
> +                                   cx_frame->have_sse ? 3 : 0, pict_type);
> +
> +    if (cx_frame->have_sse) {
> +        int i;
> +        for (i = 0; i < 3; ++i) {
> +            avctx->error[i] += cx_frame->sse[i + 1];
> +        }
> +        cx_frame->have_sse = 0;
> +    }
>
>      if (avctx->flags & AV_CODEC_FLAG_GLOBAL_HEADER) {
>          ret = av_bsf_send_packet(ctx->bsf, pkt);
> @@ -585,7 +624,7 @@ static int queue_frames(AVCodecContext *avctx,
> AVPacket *pkt_out)
>                  /* avoid storing the frame when the list is empty and we
> haven't yet
>                   * provided a frame for output */
>                  av_assert0(!ctx->coded_frame_list);
> -                cx_pktcpy(&cx_frame, pkt);
> +                cx_pktcpy(&cx_frame, pkt, ctx);
>                  size = storeframe(avctx, &cx_frame, pkt_out);
>                  if (size < 0)
>                      return size;
> @@ -598,7 +637,7 @@ static int queue_frames(AVCodecContext *avctx,
> AVPacket *pkt_out)
>                             "Frame queue element alloc failed\n");
>                      return AVERROR(ENOMEM);
>                  }
> -                cx_pktcpy(cx_frame, pkt);
> +                cx_pktcpy(cx_frame, pkt, ctx);
>                  cx_frame->buf = av_malloc(cx_frame->sz);
>
>                  if (!cx_frame->buf) {
> @@ -628,7 +667,16 @@ static int queue_frames(AVCodecContext *avctx,
> AVPacket *pkt_out)
>              stats->sz += pkt->data.twopass_stats.sz;
>              break;
>          }
> -        case AOM_CODEC_PSNR_PKT: // FIXME add support for
> AV_CODEC_FLAG_PSNR
> +        case AOM_CODEC_PSNR_PKT:
> +        {
> +            av_assert0(!ctx->have_sse);
> +            ctx->sse[0] = pkt->data.psnr.sse[0];
> +            ctx->sse[1] = pkt->data.psnr.sse[1];
> +            ctx->sse[2] = pkt->data.psnr.sse[2];
> +            ctx->sse[3] = pkt->data.psnr.sse[3];
> +            ctx->have_sse = 1;
> +            break;
> +        }
>          case AOM_CODEC_CUSTOM_PKT:
>              // ignore unsupported/unrecognized packet types
>              break;
> --
> 2.19.0.605.g01d371f741-goog
>
>
James Almer Oct. 12, 2018, 1:11 a.m. UTC | #2
On 10/11/2018 9:49 PM, Sam John wrote:
> James,
> 
> Is there anything to be done for this patch? Can you please approve this
> change and apply the patch to ffmpeg head.

Sorry for the delay. Pushed it after changing it so psnr is only
exported when every frame can be properly tagged as intra or inter.

Thanks.
diff mbox

Patch

diff --git a/libavcodec/libaomenc.c b/libavcodec/libaomenc.c
index 6a79d9b873..481338406c 100644
--- a/libavcodec/libaomenc.c
+++ b/libavcodec/libaomenc.c
@@ -50,6 +50,9 @@  struct FrameListData {
     unsigned long duration;          /**< duration to show frame
                                           (in timebase units) */
     uint32_t flags;                  /**< flags for this frame */
+    uint64_t sse[4];
+    int have_sse;                    /**< true if we have pending sse[] */
+    uint64_t frame_number;
     struct FrameListData *next;
 };
 
@@ -68,6 +71,9 @@  typedef struct AOMEncoderContext {
     int static_thresh;
     int drop_threshold;
     int noise_sensitivity;
+    uint64_t sse[4];
+    int have_sse; /**< true if we have pending sse[] */
+    uint64_t frame_number;
 } AOMContext;
 
 static const char *const ctlidstr[] = {
@@ -289,7 +295,8 @@  static av_cold int aom_init(AVCodecContext *avctx,
 {
     AOMContext *ctx = avctx->priv_data;
     struct aom_codec_enc_cfg enccfg = { 0 };
-    aom_codec_flags_t flags = 0;
+    aom_codec_flags_t flags =
+        (avctx->flags & AV_CODEC_FLAG_PSNR) ? AOM_CODEC_USE_PSNR : 0;
     AVCPBProperties *cpb_props;
     int res;
     aom_img_fmt_t img_fmt;
@@ -499,13 +506,23 @@  static av_cold int aom_init(AVCodecContext *avctx,
 }
 
 static inline void cx_pktcpy(struct FrameListData *dst,
-                             const struct aom_codec_cx_pkt *src)
+                             const struct aom_codec_cx_pkt *src,
+                             AOMContext *ctx)
 {
     dst->pts      = src->data.frame.pts;
     dst->duration = src->data.frame.duration;
     dst->flags    = src->data.frame.flags;
     dst->sz       = src->data.frame.sz;
     dst->buf      = src->data.frame.buf;
+    dst->have_sse = 0;
+    dst->frame_number = ++ctx->frame_number;
+    dst->have_sse = ctx->have_sse;
+    if (ctx->have_sse) {
+        /* associate last-seen SSE to the frame. */
+        /* Transfers ownership from ctx to dst. */
+        memcpy(dst->sse, ctx->sse, sizeof(dst->sse));
+        ctx->have_sse = 0;
+    }
 }
 
 /**
@@ -519,6 +536,7 @@  static int storeframe(AVCodecContext *avctx, struct FrameListData *cx_frame,
                       AVPacket *pkt)
 {
     AOMContext *ctx = avctx->priv_data;
+    int pict_type;
     int ret = ff_alloc_packet2(avctx, pkt, cx_frame->sz, 0);
     if (ret < 0) {
         av_log(avctx, AV_LOG_ERROR,
@@ -528,8 +546,29 @@  static int storeframe(AVCodecContext *avctx, struct FrameListData *cx_frame,
     memcpy(pkt->data, cx_frame->buf, pkt->size);
     pkt->pts = pkt->dts = cx_frame->pts;
 
-    if (!!(cx_frame->flags & AOM_FRAME_IS_KEY))
+    if (!!(cx_frame->flags & AOM_FRAME_IS_KEY)) {
+        pict_type = AV_PICTURE_TYPE_I;
         pkt->flags |= AV_PKT_FLAG_KEY;
+    } else {
+#ifdef AOM_FRAME_IS_INTRAONLY
+        if (cx_frame->flags & AOM_FRAME_IS_INTRAONLY)
+            pict_type = AV_PICTURE_TYPE_I;
+        else
+#else
+            pict_type = AV_PICTURE_TYPE_P;
+#endif
+    }
+
+    ff_side_data_set_encoder_stats(pkt, 0, cx_frame->sse + 1,
+                                   cx_frame->have_sse ? 3 : 0, pict_type);
+
+    if (cx_frame->have_sse) {
+        int i;
+        for (i = 0; i < 3; ++i) {
+            avctx->error[i] += cx_frame->sse[i + 1];
+        }
+        cx_frame->have_sse = 0;
+    }
 
     if (avctx->flags & AV_CODEC_FLAG_GLOBAL_HEADER) {
         ret = av_bsf_send_packet(ctx->bsf, pkt);
@@ -585,7 +624,7 @@  static int queue_frames(AVCodecContext *avctx, AVPacket *pkt_out)
                 /* avoid storing the frame when the list is empty and we haven't yet
                  * provided a frame for output */
                 av_assert0(!ctx->coded_frame_list);
-                cx_pktcpy(&cx_frame, pkt);
+                cx_pktcpy(&cx_frame, pkt, ctx);
                 size = storeframe(avctx, &cx_frame, pkt_out);
                 if (size < 0)
                     return size;
@@ -598,7 +637,7 @@  static int queue_frames(AVCodecContext *avctx, AVPacket *pkt_out)
                            "Frame queue element alloc failed\n");
                     return AVERROR(ENOMEM);
                 }
-                cx_pktcpy(cx_frame, pkt);
+                cx_pktcpy(cx_frame, pkt, ctx);
                 cx_frame->buf = av_malloc(cx_frame->sz);
 
                 if (!cx_frame->buf) {
@@ -628,7 +667,16 @@  static int queue_frames(AVCodecContext *avctx, AVPacket *pkt_out)
             stats->sz += pkt->data.twopass_stats.sz;
             break;
         }
-        case AOM_CODEC_PSNR_PKT: // FIXME add support for AV_CODEC_FLAG_PSNR
+        case AOM_CODEC_PSNR_PKT:
+        {
+            av_assert0(!ctx->have_sse);
+            ctx->sse[0] = pkt->data.psnr.sse[0];
+            ctx->sse[1] = pkt->data.psnr.sse[1];
+            ctx->sse[2] = pkt->data.psnr.sse[2];
+            ctx->sse[3] = pkt->data.psnr.sse[3];
+            ctx->have_sse = 1;
+            break;
+        }
         case AOM_CODEC_CUSTOM_PKT:
             // ignore unsupported/unrecognized packet types
             break;