Add custom codec callback in avifEncoder

CHANGELOG.md

@@ -29,6 +29,8 @@ The changes are relative to the previous release, unless the baseline is specifi
   avifGainMapMetadataDouble structs.
 * Add avif(Un)SignedFraction structs and avifDoubleTo(Un)SignedFraction
   utility functions.
+* Add customEncodeImageFunc, customEncodeFinishFunc and customEncodeData fields
+  to the avifEncoder struct to override the AV1 codec.
 
 ## [1.1.1] - 2024-07-30
 

include/avif/avif.h

@@ -1414,15 +1414,31 @@ AVIF_API avifResult avifDecoderNthImageMaxExtent(const avifDecoder * decoder, ui
 // ---------------------------------------------------------------------------
 // avifEncoder
 
+struct avifEncoder;
 struct avifEncoderData;
 struct avifCodecSpecificOptions;
+struct avifEncoderCustomEncodeImageArgs;
 
 typedef struct avifScalingMode
 {
     avifFraction horizontal;
     avifFraction vertical;
 } avifScalingMode;
 
+// If enabled in avifEncoder, called for each coded image item (color, alpha, and gain map), grid cell,
+// and sequence (color, alpha).
+// Returns AVIF_RESULT_OK if it overrides the AV1 codec encoding pipeline for that element.
+// Returns AVIF_RESULT_NO_CONTENT if the AV1 codec encoding pipeline should be run for that element.
+// Returns an error otherwise.
+typedef avifResult (*avifEncoderCustomEncodeImageFunc)(struct avifEncoder * encoder,
+                                                       const avifImage * image,
+                                                       const struct avifEncoderCustomEncodeImageArgs * args);
+// Only called if avifEncoderCustomEncodeImageFunc returned AVIF_RESULT_OK.
+// Returns AVIF_RESULT_OK every time it outputs an AV1 sample.
+// Returns AVIF_RESULT_NO_IMAGES_REMAINING once all samples were output.
+// Returns an error otherwise.
+typedef avifResult (*avifEncoderCustomEncodeFinishFunc)(struct avifEncoder * encoder, avifROData * sample);
+
 // Notes:
 // * The avifEncoder struct may be extended in a future release. Code outside the libavif library
 //   must allocate avifEncoder by calling the avifEncoderCreate() function.
@@ -1495,6 +1511,13 @@ typedef struct avifEncoder
 
     // Version 1.1.0 ends here. Add any new members after this line.
 
+    // Override the AV1 codec if both not null. Warning: Experimental feature.
+    // May be used to provide the payload of an AV1 coded image item or sequence.
+    avifEncoderCustomEncodeImageFunc customEncodeImageFunc;
+    avifEncoderCustomEncodeFinishFunc customEncodeFinishFunc;
+    // Ignored by libavif. May be used by customEncodeImageFunc and customEncodeFinishFunc to point to user data.
+    void * customEncodeData;
+
 #if defined(AVIF_ENABLE_EXPERIMENTAL_GAIN_MAP)
     int qualityGainMap; // changeable encoder setting
 #endif
@@ -1524,6 +1547,22 @@ typedef enum avifAddImageFlag
 } avifAddImageFlag;
 typedef uint32_t avifAddImageFlags;
 
+// Arguments passed to avifEncoderCustomEncodeImageFunc by the avifEncoder instance.
+typedef struct avifEncoderCustomEncodeImageArgs
+{
+    // Encoding settings requested by the avifEncoder instance for the current AV1 coded image item or sequence.
+    avifAddImageFlags addImageFlags;
+    int quantizer;    // AV1 quality setting in range [AVIF_QUANTIZER_BEST_QUALITY:AVIF_QUANTIZER_WORST_QUALITY].
+    int tileRowsLog2; // Logarithm in base 2 of the number of AV1 tile rows.
+    int tileColsLog2; // Logarithm in base 2 of the number of AV1 tile columns.
+
+    // Description of the current AV1 coded image item or sequence.
+    avifBool isAlpha; // True if the current AV1 image item or sequence holds the translucency layer.
+#if defined(AVIF_ENABLE_EXPERIMENTAL_GAIN_MAP)
+    avifBool isGainmap; // True if the current AV1 image item or sequence holds the HDR gainmap layer.
+#endif
+} avifEncoderCustomEncodeImageArgs;
+
 // Multi-function alternative to avifEncoderWrite() for advanced features.
 //
 // Usage / function call order is:

src/write.c

@@ -243,6 +243,8 @@ typedef struct avifEncoderData
     // Fields specific to AV1/AV2
     const char * imageItemType;  // "av01" for AV1 ("av02" for AV2 if AVIF_CODEC_AVM)
     const char * configPropName; // "av1C" for AV1 ("av2C" for AV2 if AVIF_CODEC_AVM)
+    // Custom AV1 encoding function
+    avifBool customEncodeImageFuncUsed;
 } avifEncoderData;
 
 static void avifEncoderDataDestroy(avifEncoderData * data);
@@ -2104,17 +2106,38 @@ static avifResult avifEncoderAddImageInternal(avifEncoder * encoder,
             // If alpha channel is present, set disableLaggedOutput to AVIF_TRUE. If the encoder supports it, this enables
             // avifEncoderDataShouldForceKeyframeForAlpha to force a keyframe in the alpha channel whenever a keyframe has been
             // encoded in the color channel for animated images.
-            avifResult encodeResult = item->codec->encodeImage(item->codec,
-                                                               encoder,
-                                                               cellImage,
-                                                               isAlpha,
-                                                               encoder->data->tileRowsLog2,
-                                                               encoder->data->tileColsLog2,
-                                                               quantizer,
-                                                               encoderChanges,
-                                                               /*disableLaggedOutput=*/encoder->data->alphaPresent,
-                                                               addImageFlags,
-                                                               item->encodeOutput);
+            const avifBool disableLaggedOutput = encoder->data->alphaPresent;
+
+            avifResult encodeResult = AVIF_RESULT_NO_CONTENT;
+            if (encoder->customEncodeImageFunc != NULL && encoder->customEncodeFinishFunc != NULL) {
+                const avifEncoderCustomEncodeImageArgs args = {
+                    addImageFlags,
+                    quantizer,
+                    encoder->data->tileRowsLog2,
+                    encoder->data->tileColsLog2,
+                    isAlpha,
+#if defined(AVIF_ENABLE_EXPERIMENTAL_GAIN_MAP)
+                    item->itemCategory == AVIF_ITEM_GAIN_MAP
+#endif
+                };
+                encodeResult = encoder->customEncodeImageFunc(encoder, cellImage, &args);
+                encoder->data->customEncodeImageFuncUsed = encodeResult != AVIF_RESULT_NO_CONTENT;
+            }
+
+            if (encodeResult == AVIF_RESULT_NO_CONTENT) {
+                encodeResult = item->codec->encodeImage(item->codec,
+                                                        encoder,
+                                                        cellImage,
+                                                        isAlpha,
+                                                        encoder->data->tileRowsLog2,
+                                                        encoder->data->tileColsLog2,
+                                                        quantizer,
+                                                        encoderChanges,
+                                                        disableLaggedOutput,
+                                                        addImageFlags,
+                                                        item->encodeOutput);
+            }
+
 #if defined(AVIF_ENABLE_EXPERIMENTAL_SAMPLE_TRANSFORM)
             // Revert quality settings if they changed.
             if (*encoderMinQuantizer != originalMinQuantizer || *encoderMaxQuantizer != originalMaxQuantizer) {
@@ -3094,7 +3117,14 @@ avifResult avifEncoderFinish(avifEncoder * encoder, avifRWData * output)
     for (uint32_t itemIndex = 0; itemIndex < encoder->data->items.count; ++itemIndex) {
         avifEncoderItem * item = &encoder->data->items.item[itemIndex];
         if (item->codec) {
-            if (!item->codec->encodeFinish(item->codec, item->encodeOutput)) {
+            if (encoder->data->customEncodeImageFuncUsed) {
+                avifROData sample = AVIF_DATA_EMPTY;
+                avifResult encodeResult;
+                while ((encodeResult = encoder->customEncodeFinishFunc(encoder, &sample)) != AVIF_RESULT_NO_IMAGES_REMAINING) {
+                    AVIF_CHECKRES(encodeResult);
+                    AVIF_CHECKRES(avifCodecEncodeOutputAddSample(item->encodeOutput, sample.data, sample.size, /*sync=*/AVIF_TRUE));
+                }
+            } else if (!item->codec->encodeFinish(item->codec, item->encodeOutput)) {
                 return avifGetErrorForItemCategory(item->itemCategory);
             }
 

tests/CMakeLists.txt

@@ -109,6 +109,7 @@ if(AVIF_ENABLE_GTEST)
     add_avif_gtest(avifcodectest)
     add_avif_internal_gtest_with_data(avifcolrconverttest)
     add_avif_internal_gtest(avifcolrtest)
+    add_avif_gtest(avifcustomtest)
     add_avif_gtest_with_data(avifdecodetest)
     add_avif_gtest_with_data(avifdimgtest avifincrtest_helpers)
     add_avif_gtest_with_data(avifencodetest)
@@ -354,6 +355,7 @@ if(AVIF_CODEC_AVM_ENABLED)
                 avifchangesettingtest
                 avifcllitest
                 avifcolrconverttest
+                avifcustomtest
                 avifdimgtest
                 avifencodetest
                 avifgridapitest

tests/gtest/avifcustomtest.cc

@@ -0,0 +1,70 @@
+// Copyright 2024 Google LLC
+// SPDX-License-Identifier: BSD-2-Clause
+
+#include <algorithm>
+
+#include "avif/avif.h"
+#include "aviftest_helpers.h"
+#include "gtest/gtest.h"
+
+namespace avif {
+namespace {
+
+avifResult CustomEncodeImageFunc(avifEncoder* encoder, const avifImage*,
+                                 const avifEncoderCustomEncodeImageArgs*) {
+  if (encoder->customEncodeData != NULL) {
+    return AVIF_RESULT_OK;  // Overrides the AV1 codec encoding pipeline.
+  } else {
+    return AVIF_RESULT_NO_CONTENT;  // Lets libavif encode the image item.
+  }
+}
+
+avifResult CustomEncodeFinishFunc(avifEncoder* encoder, avifROData* sample) {
+  avifROData* av1_payload =
+      reinterpret_cast<avifROData*>(encoder->customEncodeData);
+  if (av1_payload->size != 0) {
+    *sample = *av1_payload;
+    *av1_payload = AVIF_DATA_EMPTY;
+    return AVIF_RESULT_OK;  // Outputs a sample.
+  } else {
+    return AVIF_RESULT_NO_IMAGES_REMAINING;  // Done.
+  }
+}
+
+TEST(BasicTest, EncodeDecode) {
+  ImagePtr image = testutil::CreateImage(12, 34, 8, AVIF_PIXEL_FORMAT_YUV420,
+                                         AVIF_PLANES_YUV);
+  ASSERT_NE(image, nullptr);
+  testutil::FillImageGradient(image.get());
+
+  EncoderPtr encoder(avifEncoderCreate());
+  ASSERT_NE(encoder, nullptr);
+  testutil::AvifRwData encoded;
+  ASSERT_EQ(avifEncoderWrite(encoder.get(), image.get(), &encoded),
+            AVIF_RESULT_OK);
+
+  const uint8_t* kMdat = reinterpret_cast<const uint8_t*>("mdat");
+  const uint8_t* mdat_position =
+      std::search(encoded.data, encoded.data + encoded.size, kMdat, kMdat + 4);
+  ASSERT_NE(mdat_position, encoded.data + encoded.size);
+  avifROData av1_payload{
+      mdat_position + 4,
+      static_cast<size_t>((encoded.data + encoded.size) - (mdat_position + 4))};
+
+  EncoderPtr encoder_custom(avifEncoderCreate());
+  ASSERT_NE(encoder_custom, nullptr);
+  encoder_custom->customEncodeData = reinterpret_cast<void*>(&av1_payload);
+  encoder_custom->customEncodeImageFunc = CustomEncodeImageFunc;
+  encoder_custom->customEncodeFinishFunc = CustomEncodeFinishFunc;
+  testutil::AvifRwData encoded_custom;
+  ASSERT_EQ(
+      avifEncoderWrite(encoder_custom.get(), image.get(), &encoded_custom),
+      AVIF_RESULT_OK);
+
+  ASSERT_EQ(encoded.size, encoded_custom.size);
+  EXPECT_TRUE(std::equal(encoded.data, encoded.data + encoded.size,
+                         encoded_custom.data));
+}
+
+}  // namespace
+}  // namespace avif