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 20180927014227.38058-1-samjohn@google.com
State Superseded
Headers show

Commit Message

Sam John Sept. 27, 2018, 1:42 a.m. UTC
---
 libavcodec/libaomenc.c | 83 ++++++++++++++++++++++++++++++++++++++----
 1 file changed, 75 insertions(+), 8 deletions(-)

Comments

James Almer Sept. 27, 2018, 2:40 a.m. UTC | #1
On 9/26/2018 10:42 PM, Sam John wrote:
> ---
>  libavcodec/libaomenc.c | 83 ++++++++++++++++++++++++++++++++++++++----
>  1 file changed, 75 insertions(+), 8 deletions(-)
> 
> diff --git a/libavcodec/libaomenc.c b/libavcodec/libaomenc.c
> index 6a79d9b873..4e75ca0b18 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,
> @@ -527,11 +545,51 @@ 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))
> +FF_DISABLE_DEPRECATION_WARNINGS
> +    avctx->coded_frame->pts       = cx_frame->pts;
> +    avctx->coded_frame->key_frame = !!(cx_frame->flags & AOM_FRAME_IS_KEY);
> +FF_ENABLE_DEPRECATION_WARNINGS

Again please, don't add deprecated code. avctx->coded_frame is not meant
to be used in new modules, and is present in old modules only until
removal for backwards compatibility.

> +
> +    if (!!(cx_frame->flags & AOM_FRAME_IS_KEY)) {
> +        pict_type = AV_PICTURE_TYPE_I;
> +FF_DISABLE_DEPRECATION_WARNINGS
> +        avctx->coded_frame->pict_type = pict_type;
> +FF_ENABLE_DEPRECATION_WARNINGS
>          pkt->flags |= AV_PKT_FLAG_KEY;
> +    } else {
> +#ifdef AOM_FRAME_IS_INTRAONLY

Nice, glad to see this was implemented.

If no one else is against it I'm fine keeping the #else case for the
time being even if it may flag intra frames as inter in the encoder
stats since i plan on bumping the minimum required version as soon as a
new libaom release is tagged. A lot of development has taken place and
plenty of API was added that makes 1.0.0 unoptimal.

> +        if (cx_frame->flags & AOM_FRAME_IS_INTRAONLY)
> +            pict_type = AV_PICTURE_TYPE_I;
> +        else
> +            pict_type = AV_PICTURE_TYPE_P;
> +#else

Nit: You can reduce duplication by removing the two lines below and
moving the #else one line up.

> +        pict_type = AV_PICTURE_TYPE_P;
> +#endif
> +
> +FF_DISABLE_DEPRECATION_WARNINGS
> +        avctx->coded_frame->pict_type = pict_type;
> +FF_ENABLE_DEPRECATION_WARNINGS
> +    }
>  
> -    if (avctx->flags & AV_CODEC_FLAG_GLOBAL_HEADER) {
> +    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;
> +        /* Beware of the Y/U/V/all order! */
> +FF_DISABLE_DEPRECATION_WARNINGS
> +        avctx->coded_frame->error[0] = cx_frame->sse[1];
> +        avctx->coded_frame->error[1] = cx_frame->sse[2];
> +        avctx->coded_frame->error[2] = cx_frame->sse[3];
> +        avctx->coded_frame->error[3] = 0;    // alpha
> +FF_ENABLE_DEPRECATION_WARNINGS
> +        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);
>          if (ret < 0) {
>              av_log(avctx, AV_LOG_ERROR, "extract_extradata filter "
> @@ -585,7 +643,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 +656,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 +686,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;
>
diff mbox

Patch

diff --git a/libavcodec/libaomenc.c b/libavcodec/libaomenc.c
index 6a79d9b873..4e75ca0b18 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,
@@ -527,11 +545,51 @@  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))
+FF_DISABLE_DEPRECATION_WARNINGS
+    avctx->coded_frame->pts       = cx_frame->pts;
+    avctx->coded_frame->key_frame = !!(cx_frame->flags & AOM_FRAME_IS_KEY);
+FF_ENABLE_DEPRECATION_WARNINGS
+
+    if (!!(cx_frame->flags & AOM_FRAME_IS_KEY)) {
+        pict_type = AV_PICTURE_TYPE_I;
+FF_DISABLE_DEPRECATION_WARNINGS
+        avctx->coded_frame->pict_type = pict_type;
+FF_ENABLE_DEPRECATION_WARNINGS
         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
+            pict_type = AV_PICTURE_TYPE_P;
+#else
+        pict_type = AV_PICTURE_TYPE_P;
+#endif
+
+FF_DISABLE_DEPRECATION_WARNINGS
+        avctx->coded_frame->pict_type = pict_type;
+FF_ENABLE_DEPRECATION_WARNINGS
+    }
 
-    if (avctx->flags & AV_CODEC_FLAG_GLOBAL_HEADER) {
+    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;
+        /* Beware of the Y/U/V/all order! */
+FF_DISABLE_DEPRECATION_WARNINGS
+        avctx->coded_frame->error[0] = cx_frame->sse[1];
+        avctx->coded_frame->error[1] = cx_frame->sse[2];
+        avctx->coded_frame->error[2] = cx_frame->sse[3];
+        avctx->coded_frame->error[3] = 0;    // alpha
+FF_ENABLE_DEPRECATION_WARNINGS
+        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);
         if (ret < 0) {
             av_log(avctx, AV_LOG_ERROR, "extract_extradata filter "
@@ -585,7 +643,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 +656,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 +686,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;