lavc: add API for exporting reconstructed frames from encoders

11 files changed, 87 insertions, 8 deletions

diff --git a/doc/APIchanges b/doc/APIchanges
index f13c88dc50..0e9ea4d7c5 100644
--- a/doc/APIchanges
+++ b/doc/APIchanges
@@ -14,6 +14,11 @@ libavutil:     2021-04-27
 
 API changes, most recent first:
 
+2022-08-xx - xxxxxxxxxx - lavc 59.41.100 - avcodec.h codec.h
+  Add AV_CODEC_FLAG_RECON_FRAME and AV_CODEC_CAP_ENCODER_RECON_FRAME.
+  avcodec_receive_frame() may now be used on encoders when
+  AV_CODEC_FLAG_RECON_FRAME is active.
+
 2022-08-xx - xxxxxxxxxx - lavu 57.31.100 - frame.h
   av_frame_make_writable() may now be called on non-refcounted
   frames and will make a refcounted copy out of them.
diff --git a/libavcodec/avcodec.c b/libavcodec/avcodec.c
index c9105c5df2..f82d9e9f74 100644
--- a/libavcodec/avcodec.c
+++ b/libavcodec/avcodec.c
@@ -388,6 +388,8 @@ void avcodec_flush_buffers(AVCodecContext *avctx)
         }
         if (avci->in_frame)
             av_frame_unref(avci->in_frame);
+        if (avci->recon_frame)
+            av_frame_unref(avci->recon_frame);
     } else {
         av_packet_unref(avci->last_pkt_props);
         while (av_fifo_read(avci->pkt_props, avci->last_pkt_props, 1) >= 0)
@@ -468,6 +470,7 @@ av_cold int avcodec_close(AVCodecContext *avctx)
 
         av_packet_free(&avci->in_pkt);
         av_frame_free(&avci->in_frame);
+        av_frame_free(&avci->recon_frame);
 
         av_buffer_unref(&avci->pool);
 
@@ -718,3 +721,12 @@ int avcodec_is_open(AVCodecContext *s)
 {
     return !!s->internal;
 }
+
+int attribute_align_arg avcodec_receive_frame(AVCodecContext *avctx, AVFrame *frame)
+{
+    av_frame_unref(frame);
+
+    if (av_codec_is_decoder(avctx->codec))
+        return ff_decode_receive_frame(avctx, frame);
+    return ff_encode_receive_frame(avctx, frame);
+}
diff --git a/libavcodec/avcodec.h b/libavcodec/avcodec.h
index 60b215d2e9..bfd00ec493 100644
--- a/libavcodec/avcodec.h
+++ b/libavcodec/avcodec.h
@@ -232,6 +232,16 @@ typedef struct RcOverride{
  */
 #define AV_CODEC_FLAG_DROPCHANGED     (1 <<  5)
 /**
+ * Request the encoder to output reconstructed frames, i.e. frames that would be
+ * produced by decoding the encoded bistream. These frames may be retrieved by
+ * calling avcodec_receive_frame() immediately after a successful call to
+ * avcodec_receive_packet().
+ *
+ * Should only be used with encoders flagged with the
+ * AV_CODEC_CAP_ENCODER_RECON_FRAME capability.
+ */
+#define AV_CODEC_FLAG_RECON_FRAME     (1 <<  6)
+/**
  * Use internal 2pass ratecontrol in first pass mode.
  */
 #define AV_CODEC_FLAG_PASS1           (1 <<  9)
@@ -2601,21 +2611,23 @@ int avcodec_decode_subtitle2(AVCodecContext *avctx, AVSubtitle *sub,
 int avcodec_send_packet(AVCodecContext *avctx, const AVPacket *avpkt);
 
 /**
- * Return decoded output data from a decoder.
+ * Return decoded output data from a decoder or encoder (when the
+ * AV_CODEC_FLAG_RECON_FRAME flag is used).
  *
  * @param avctx codec context
  * @param frame This will be set to a reference-counted video or audio
  *              frame (depending on the decoder type) allocated by the
- *              decoder. Note that the function will always call
+ *              codec. Note that the function will always call
  *              av_frame_unref(frame) before doing anything else.
  *
  * @return
  *      0:                 success, a frame was returned
  *      AVERROR(EAGAIN):   output is not available in this state - user must try
  *                         to send new input
- *      AVERROR_EOF:       the decoder has been fully flushed, and there will be
+ *      AVERROR_EOF:       the codec has been fully flushed, and there will be
  *                         no more output frames
- *      AVERROR(EINVAL):   codec not opened, or it is an encoder
+ *      AVERROR(EINVAL):   codec not opened, or it is an encoder without
+ *                         the AV_CODEC_FLAG_RECON_FRAME flag enabled
  *      AVERROR_INPUT_CHANGED:   current decoded frame has changed parameters
  *                               with respect to first decoded frame. Applicable
  *                               when flag AV_CODEC_FLAG_DROPCHANGED is set.
diff --git a/libavcodec/codec.h b/libavcodec/codec.h
index 03e8be90a2..77a1a3f5a2 100644
--- a/libavcodec/codec.h
+++ b/libavcodec/codec.h
@@ -183,6 +183,14 @@
 #define AV_CODEC_CAP_ENCODER_FLUSH   (1 << 21)
 
 /**
+ * The encoder is able to output reconstructed frame data, i.e. raw frames that
+ * would be produced by decoding the encoded bitstream.
+ *
+ * Reconstructed frame output is enabled by the AV_CODEC_FLAG_RECON_FRAME flag.
+ */
+#define AV_CODEC_CAP_ENCODER_RECON_FRAME (1 << 22)
+
+/**
  * AVProfile.
  */
 typedef struct AVProfile {
diff --git a/libavcodec/decode.c b/libavcodec/decode.c
index 92958745df..75373989c6 100644
--- a/libavcodec/decode.c
+++ b/libavcodec/decode.c
@@ -697,13 +697,11 @@ static int apply_cropping(AVCodecContext *avctx, AVFrame *frame)
                                           AV_FRAME_CROP_UNALIGNED : 0);
 }
 
-int attribute_align_arg avcodec_receive_frame(AVCodecContext *avctx, AVFrame *frame)
+int ff_decode_receive_frame(AVCodecContext *avctx, AVFrame *frame)
 {
     AVCodecInternal *avci = avctx->internal;
     int ret, changed;
 
-    av_frame_unref(frame);
-
     if (!avcodec_is_open(avctx) || !av_codec_is_decoder(avctx->codec))
         return AVERROR(EINVAL);
 
diff --git a/libavcodec/decode.h b/libavcodec/decode.h
index 1b40f714e1..f7828fbff4 100644
--- a/libavcodec/decode.h
+++ b/libavcodec/decode.h
@@ -54,6 +54,11 @@ typedef struct FrameDecodeData {
 } FrameDecodeData;
 
 /**
+ * avcodec_receive_frame() implementation for decoders.
+ */
+int ff_decode_receive_frame(AVCodecContext *avctx, AVFrame *frame);
+
+/**
  * Called by decoders to get the next packet for decoding.
  *
  * @param pkt An empty packet to be filled with data.
diff --git a/libavcodec/encode.c b/libavcodec/encode.c
index 7cf13bf6d6..7919b165da 100644
--- a/libavcodec/encode.c
+++ b/libavcodec/encode.c
@@ -656,6 +656,18 @@ int ff_encode_preinit(AVCodecContext *avctx)
             return AVERROR(ENOMEM);
     }
 
+    if ((avctx->flags & AV_CODEC_FLAG_RECON_FRAME)) {
+        if (!(avctx->codec->capabilities & AV_CODEC_CAP_ENCODER_RECON_FRAME)) {
+            av_log(avctx, AV_LOG_ERROR, "Reconstructed frame output requested "
+                   "from an encoder not supporting it\n");
+            return AVERROR(ENOSYS);
+        }
+
+        avci->recon_frame = av_frame_alloc();
+        if (!avci->recon_frame)
+            return AVERROR(ENOMEM);
+    }
+
     return 0;
 }
 
@@ -692,3 +704,16 @@ int ff_encode_alloc_frame(AVCodecContext *avctx, AVFrame *frame)
 
     return 0;
 }
+
+int ff_encode_receive_frame(AVCodecContext *avctx, AVFrame *frame)
+{
+    AVCodecInternal *avci = avctx->internal;
+
+    if (!avci->recon_frame)
+        return AVERROR(EINVAL);
+    if (!avci->recon_frame->buf[0])
+        return avci->draining_done ? AVERROR_EOF : AVERROR(EAGAIN);
+
+    av_frame_move_ref(frame, avci->recon_frame);
+    return 0;
+}
diff --git a/libavcodec/encode.h b/libavcodec/encode.h
index b2536bf0f3..bc77918d8f 100644
--- a/libavcodec/encode.h
+++ b/libavcodec/encode.h
@@ -27,6 +27,11 @@
 #include "packet.h"
 
 /**
+ * avcodec_receive_frame() implementation for encoders.
+ */
+int ff_encode_receive_frame(AVCodecContext *avctx, AVFrame *frame);
+
+/**
  * Called by encoders to get the next frame for encoding.
  *
  * @param frame An empty frame to be filled with data.
diff --git a/libavcodec/internal.h b/libavcodec/internal.h
index 8809a7079a..8e5e2a0f59 100644
--- a/libavcodec/internal.h
+++ b/libavcodec/internal.h
@@ -108,6 +108,14 @@ typedef struct AVCodecInternal {
     AVFrame *in_frame;
 
     /**
+     * When the AV_CODEC_FLAG_RECON_FRAME flag is used. the encoder should store
+     * here the reconstructed frame corresponding to the last returned packet.
+     *
+     * Not allocated in other cases.
+     */
+    AVFrame *recon_frame;
+
+    /**
      * If this is set, then FFCodec->close (if existing) needs to be called
      * for the parent AVCodecContext.
      */
diff --git a/libavcodec/options_table.h b/libavcodec/options_table.h
index a72085ac90..cd02f5096f 100644
--- a/libavcodec/options_table.h
+++ b/libavcodec/options_table.h
@@ -57,6 +57,7 @@ static const AVOption avcodec_options[] = {
 {"qpel", "use 1/4-pel motion compensation", 0, AV_OPT_TYPE_CONST, {.i64 = AV_CODEC_FLAG_QPEL }, INT_MIN, INT_MAX, V|E, "flags"},
 {"loop", "use loop filter", 0, AV_OPT_TYPE_CONST, {.i64 = AV_CODEC_FLAG_LOOP_FILTER }, INT_MIN, INT_MAX, V|E, "flags"},
 {"qscale", "use fixed qscale", 0, AV_OPT_TYPE_CONST, {.i64 = AV_CODEC_FLAG_QSCALE }, INT_MIN, INT_MAX, 0, "flags"},
+{"recon_frame", "export reconstructed frames", 0, AV_OPT_TYPE_CONST, {.i64 = AV_CODEC_FLAG_RECON_FRAME}, .unit = "flags"},
 {"pass1", "use internal 2-pass ratecontrol in first  pass mode", 0, AV_OPT_TYPE_CONST, {.i64 = AV_CODEC_FLAG_PASS1 }, INT_MIN, INT_MAX, 0, "flags"},
 {"pass2", "use internal 2-pass ratecontrol in second pass mode", 0, AV_OPT_TYPE_CONST, {.i64 = AV_CODEC_FLAG_PASS2 }, INT_MIN, INT_MAX, 0, "flags"},
 {"gray", "only decode/encode grayscale", 0, AV_OPT_TYPE_CONST, {.i64 = AV_CODEC_FLAG_GRAY }, INT_MIN, INT_MAX, V|E|D, "flags"},
diff --git a/libavcodec/version.h b/libavcodec/version.h
index 19f3f4a272..4a1205c718 100644
--- a/libavcodec/version.h
+++ b/libavcodec/version.h
@@ -29,7 +29,7 @@
 
 #include "version_major.h"
 
-#define LIBAVCODEC_VERSION_MINOR  40
+#define LIBAVCODEC_VERSION_MINOR  41
 #define LIBAVCODEC_VERSION_MICRO 100
 
 #define LIBAVCODEC_VERSION_INT  AV_VERSION_INT(LIBAVCODEC_VERSION_MAJOR, \