avcodec/libaomenc: support AV_CODEC_CAP_ENCODER_RECON_FRAME

Signed-off-by: James Almer <jamrial@gmail.com>

5 files changed, 183 insertions, 27 deletions

diff --git a/libavcodec/Makefile b/libavcodec/Makefile
index aff7752856..ca5a38bee1 100644
--- a/libavcodec/Makefile
+++ b/libavcodec/Makefile
@@ -1055,8 +1055,8 @@ OBJS-$(CONFIG_ALAC_AT_ENCODER)            += audiotoolboxenc.o
 OBJS-$(CONFIG_ILBC_AT_ENCODER)            += audiotoolboxenc.o
 OBJS-$(CONFIG_PCM_ALAW_AT_ENCODER)        += audiotoolboxenc.o
 OBJS-$(CONFIG_PCM_MULAW_AT_ENCODER)       += audiotoolboxenc.o
-OBJS-$(CONFIG_LIBAOM_AV1_DECODER)         += libaomdec.o
-OBJS-$(CONFIG_LIBAOM_AV1_ENCODER)         += libaomenc.o
+OBJS-$(CONFIG_LIBAOM_AV1_DECODER)         += libaomdec.o libaom.o
+OBJS-$(CONFIG_LIBAOM_AV1_ENCODER)         += libaomenc.o libaom.o
 OBJS-$(CONFIG_LIBARIBB24_DECODER)         += libaribb24.o ass.o
 OBJS-$(CONFIG_LIBCELT_DECODER)            += libcelt_dec.o
 OBJS-$(CONFIG_LIBCODEC2_DECODER)          += libcodec2.o
diff --git a/libavcodec/libaom.c b/libavcodec/libaom.c
new file mode 100644
index 0000000000..0befaaa530
--- /dev/null
+++ b/libavcodec/libaom.c
@@ -0,0 +1,49 @@
+/*
+ * This file is part of FFmpeg.
+ *
+ * FFmpeg is free software; you can redistribute it and/or
+ * modify it under the terms of the GNU Lesser General Public
+ * License as published by the Free Software Foundation; either
+ * version 2.1 of the License, or (at your option) any later version.
+ *
+ * FFmpeg is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
+ * Lesser General Public License for more details.
+ *
+ * You should have received a copy of the GNU Lesser General Public
+ * License along with FFmpeg; if not, write to the Free Software
+ * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
+ */
+
+/**
+ * @file
+ * AOM common functions
+ */
+
+#include "libavutil/pixdesc.h"
+#include "libaom.h"
+
+void ff_aom_image_copy_16_to_8(AVFrame *pic, struct aom_image *img)
+{
+    const AVPixFmtDescriptor *desc = av_pix_fmt_desc_get(pic->format);
+    int i;
+
+    for (i = 0; i < desc->nb_components; i++) {
+        int w = img->d_w;
+        int h = img->d_h;
+        int x, y;
+
+        if (i) {
+            w = (w + img->x_chroma_shift) >> img->x_chroma_shift;
+            h = (h + img->y_chroma_shift) >> img->y_chroma_shift;
+        }
+
+        for (y = 0; y < h; y++) {
+            uint16_t *src = (uint16_t *)(img->planes[i] + y * img->stride[i]);
+            uint8_t *dst = pic->data[i] + y * pic->linesize[i];
+            for (x = 0; x < w; x++)
+                *dst++ = *src++;
+        }
+    }
+}
diff --git a/libavcodec/libaom.h b/libavcodec/libaom.h
new file mode 100644
index 0000000000..bb495af982
--- /dev/null
+++ b/libavcodec/libaom.h
@@ -0,0 +1,33 @@
+/*
+ * This file is part of FFmpeg.
+ *
+ * FFmpeg is free software; you can redistribute it and/or
+ * modify it under the terms of the GNU Lesser General Public
+ * License as published by the Free Software Foundation; either
+ * version 2.1 of the License, or (at your option) any later version.
+ *
+ * FFmpeg is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
+ * Lesser General Public License for more details.
+ *
+ * You should have received a copy of the GNU Lesser General Public
+ * License along with FFmpeg; if not, write to the Free Software
+ * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
+ */
+
+/**
+ * @file
+ * AOM common functions
+ */
+
+#ifndef AVCODEC_LIBAOM_H
+#define AVCODEC_LIBAOM_H
+
+#include <aom/aom_image.h>
+
+#include "libavutil/frame.h"
+
+void ff_aom_image_copy_16_to_8(AVFrame *pic, struct aom_image *img);
+
+#endif /* AVCODEC_LIBAOM_H */
diff --git a/libavcodec/libaomdec.c b/libavcodec/libaomdec.c
index 3243610304..cb672b0e65 100644
--- a/libavcodec/libaomdec.c
+++ b/libavcodec/libaomdec.c
@@ -33,6 +33,7 @@
 #include "avcodec.h"
 #include "codec_internal.h"
 #include "internal.h"
+#include "libaom.h"
 #include "profiles.h"
 
 typedef struct AV1DecodeContext {
@@ -60,30 +61,6 @@ static av_cold int aom_init(AVCodecContext *avctx,
     return 0;
 }
 
-static void image_copy_16_to_8(AVFrame *pic, struct aom_image *img)
-{
-    const AVPixFmtDescriptor *desc = av_pix_fmt_desc_get(pic->format);
-    int i;
-
-    for (i = 0; i < desc->nb_components; i++) {
-        int w = img->d_w;
-        int h = img->d_h;
-        int x, y;
-
-        if (i) {
-            w = (w + img->x_chroma_shift) >> img->x_chroma_shift;
-            h = (h + img->y_chroma_shift) >> img->y_chroma_shift;
-        }
-
-        for (y = 0; y < h; y++) {
-            uint16_t *src = (uint16_t *)(img->planes[i] + y * img->stride[i]);
-            uint8_t *dst = pic->data[i] + y * pic->linesize[i];
-            for (x = 0; x < w; x++)
-                *dst++ = *src++;
-        }
-    }
-}
-
 // returns 0 on success, AVERROR_INVALIDDATA otherwise
 static int set_pix_fmt(AVCodecContext *avctx, struct aom_image *img)
 {
@@ -223,7 +200,7 @@ static int aom_decode(AVCodecContext *avctx, AVFrame *picture,
         ff_set_sar(avctx, picture->sample_aspect_ratio);
 
         if ((img->fmt & AOM_IMG_FMT_HIGHBITDEPTH) && img->bit_depth == 8)
-            image_copy_16_to_8(picture, img);
+            ff_aom_image_copy_16_to_8(picture, img);
         else {
             const uint8_t *planes[4] = { img->planes[0], img->planes[1], img->planes[2] };
             const int      stride[4] = { img->stride[0], img->stride[1], img->stride[2] };
diff --git a/libavcodec/libaomenc.c b/libavcodec/libaomenc.c
index b792226744..1fd69d59a7 100644
--- a/libavcodec/libaomenc.c
+++ b/libavcodec/libaomenc.c
@@ -31,6 +31,7 @@
 #include "libavutil/base64.h"
 #include "libavutil/common.h"
 #include "libavutil/cpu.h"
+#include "libavutil/imgutils.h"
 #include "libavutil/mathematics.h"
 #include "libavutil/opt.h"
 #include "libavutil/pixdesc.h"
@@ -41,6 +42,7 @@
 #include "codec_internal.h"
 #include "encode.h"
 #include "internal.h"
+#include "libaom.h"
 #include "packet_internal.h"
 #include "profiles.h"
 
@@ -208,6 +210,7 @@ static const char *const ctlidstr[] = {
 #ifdef AOM_CTRL_AV1E_GET_TARGET_SEQ_LEVEL_IDX
     [AV1E_GET_TARGET_SEQ_LEVEL_IDX]     = "AV1E_GET_TARGET_SEQ_LEVEL_IDX",
 #endif
+    [AV1_GET_NEW_FRAME_IMAGE]           = "AV1_GET_NEW_FRAME_IMAGE",
 };
 
 static av_cold void log_encoder_error(AVCodecContext *avctx, const char *desc)
@@ -364,6 +367,31 @@ static av_cold int codecctl_intp(AVCodecContext *avctx,
 }
 #endif
 
+static av_cold int codecctl_imgp(AVCodecContext *avctx,
+#ifdef UENUM1BYTE
+                                 aome_enc_control_id id,
+#else
+                                 enum aome_enc_control_id id,
+#endif
+                                 struct aom_image *img)
+{
+    AOMContext *ctx = avctx->priv_data;
+    char buf[80];
+    int res;
+
+    snprintf(buf, sizeof(buf), "%s:", ctlidstr[id]);
+
+    res = aom_codec_control(&ctx->encoder, id, img);
+    if (res != AOM_CODEC_OK) {
+        snprintf(buf, sizeof(buf), "Failed to get %s codec control",
+                 ctlidstr[id]);
+        log_encoder_error(avctx, buf);
+        return AVERROR(EINVAL);
+    }
+
+    return 0;
+}
+
 static av_cold int aom_free(AVCodecContext *avctx)
 {
     AOMContext *ctx = avctx->priv_data;
@@ -1206,6 +1234,37 @@ static int queue_frames(AVCodecContext *avctx, AVPacket *pkt_out)
     return size;
 }
 
+static enum AVPixelFormat aomfmt_to_pixfmt(struct aom_image *img)
+{
+    switch (img->fmt) {
+    case AOM_IMG_FMT_I420:
+    case AOM_IMG_FMT_I42016:
+        if (img->bit_depth == 8)
+            return img->monochrome ? AV_PIX_FMT_GRAY8 : AV_PIX_FMT_YUV420P;
+        else if (img->bit_depth == 10)
+            return img->monochrome ? AV_PIX_FMT_GRAY10 : AV_PIX_FMT_YUV420P10;
+        else
+            return img->monochrome ? AV_PIX_FMT_GRAY12 : AV_PIX_FMT_YUV420P12;
+    case AOM_IMG_FMT_I422:
+    case AOM_IMG_FMT_I42216:
+        if (img->bit_depth == 8)
+            return AV_PIX_FMT_YUV422P;
+        else if (img->bit_depth == 10)
+            return AV_PIX_FMT_YUV422P10;
+        else
+            return AV_PIX_FMT_YUV422P12;
+    case AOM_IMG_FMT_I444:
+    case AOM_IMG_FMT_I44416:
+        if (img->bit_depth == 8)
+            return AV_PIX_FMT_YUV444P;
+        else if (img->bit_depth == 10)
+            return AV_PIX_FMT_YUV444P10;
+        else
+            return AV_PIX_FMT_YUV444P12;
+    };
+    return AV_PIX_FMT_NONE;
+}
+
 static int aom_encode(AVCodecContext *avctx, AVPacket *pkt,
                       const AVFrame *frame, int *got_packet)
 {
@@ -1259,6 +1318,43 @@ static int aom_encode(AVCodecContext *avctx, AVPacket *pkt,
     }
 
     *got_packet = !!coded_size;
+
+    if (*got_packet && avctx->flags & AV_CODEC_FLAG_RECON_FRAME) {
+        AVCodecInternal *avci = avctx->internal;
+        struct aom_image img;
+
+        av_frame_unref(avci->recon_frame);
+
+        res = codecctl_imgp(avctx, AV1_GET_NEW_FRAME_IMAGE, &img);
+        if (res < 0)
+            return res;
+
+        avci->recon_frame->format = aomfmt_to_pixfmt(&img);
+        if (avci->recon_frame->format == AV_PIX_FMT_NONE) {
+            av_log(ctx, AV_LOG_ERROR,
+                   "Unhandled reconstructed frame colorspace: %d\n",
+                   img.fmt);
+            return AVERROR(ENOSYS);
+        }
+
+        avci->recon_frame->width  = img.d_w;
+        avci->recon_frame->height = img.d_h;
+
+        res = av_frame_get_buffer(avci->recon_frame, 0);
+        if (res < 0)
+            return res;
+
+        if ((img.fmt & AOM_IMG_FMT_HIGHBITDEPTH) && img.bit_depth == 8)
+            ff_aom_image_copy_16_to_8(avci->recon_frame, &img);
+        else {
+            const uint8_t *planes[4] = { img.planes[0], img.planes[1], img.planes[2] };
+            const int      stride[4] = { img.stride[0], img.stride[1], img.stride[2] };
+
+            av_image_copy(avci->recon_frame->data, avci->recon_frame->linesize, planes,
+                          stride, avci->recon_frame->format, img.d_w, img.d_h);
+        }
+    }
+
     return 0;
 }
 
@@ -1434,6 +1530,7 @@ FFCodec ff_libaom_av1_encoder = {
     .p.type         = AVMEDIA_TYPE_VIDEO,
     .p.id           = AV_CODEC_ID_AV1,
     .p.capabilities = AV_CODEC_CAP_DR1 | AV_CODEC_CAP_DELAY |
+                      AV_CODEC_CAP_ENCODER_RECON_FRAME |
                       AV_CODEC_CAP_OTHER_THREADS,
     .p.profiles     = NULL_IF_CONFIG_SMALL(ff_av1_profiles),
     .p.priv_class   = &class_aom,