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

@@ -1411,6 +1411,70 @@ typedef struct avifExtent
 // This function may be used after a successful call (AVIF_RESULT_OK) to avifDecoderParse().
 AVIF_API avifResult avifDecoderNthImageMaxExtent(const avifDecoder * decoder, uint32_t frameIndex, avifExtent * outExtent);
 
+// ---------------------------------------------------------------------------
+// Custom codec callbacks
+
+struct avifEncoder;
+typedef uint32_t avifAddImageFlags; // avifAddImageFlag bit mask.
+
+// The meaning of the pixels of the coded image item being encoded.
+typedef enum avifEncoderCustomEncodeImageItemType
+{
+    AVIF_ENCODER_CUSTOM_ENCODE_ITEM_COLOR,
+    AVIF_ENCODER_CUSTOM_ENCODE_ITEM_ALPHA,
+#if defined(AVIF_ENABLE_EXPERIMENTAL_GAIN_MAP)
+    AVIF_ENCODER_CUSTOM_ENCODE_ITEM_GAINMAP,
+#endif
+    // Other values may be valid but not exposed here.
+} avifEncoderCustomEncodeImageItemType;
+
+// Uniquely identifies the item being encoded between avifEncoderCustomEncodeImageFunc and
+// avifEncoderCustomEncodeFinishFunc calls. It can be a whole frame or a HEIF grid cell.
+typedef struct avifEncoderCustomEncodeImageItem
+{
+    avifEncoderCustomEncodeImageItemType type;
+    uint32_t id;         // 1-based image item id.
+    uint32_t gridRow;    // Vertical coordinate of this cell in its grid. 0 if there is no grid.
+    uint32_t gridColumn; // Horizontal coordinate of this cell in its grid. 0 if there is no grid.
+} avifEncoderCustomEncodeImageItem;
+
+// Arguments passed to avifEncoderCustomEncodeImageFunc corresponding to the encoding settings
+// requested by the avifEncoder instance for the current AV1 coded image item, layer, or sequence frame.
+typedef struct avifEncoderCustomEncodeImageArgs
+{
+    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.
+} avifEncoderCustomEncodeImageArgs;
+
+// If enabled in avifEncoder, called once for each coded image item (color, alpha, or gain map).
+// Further calls with the same avifEncoderCustomEncodeImageItem correspond to the following frames
+// of the sequence whose first frame is also used for that coded image item, or to the other layers
+// of that layered coded image item.
+// Returns AVIF_RESULT_OK if it overrides the AV1 codec encoding pipeline for that item (and track if any).
+// Returns AVIF_RESULT_NO_CONTENT if the AV1 codec encoding pipeline should be run.
+// Returns an error otherwise.
+// All calls with the same avifEncoderCustomEncodeImageItem must return the same status, or an error.
+// Calls to avifEncoderCustomEncodeImageFunc and avifEncoderCustomEncodeFinishFunc with the same
+// avifEncoderCustomEncodeImageItem will happen sequentially.
+// Calls to avifEncoderCustomEncodeImageFunc and avifEncoderCustomEncodeFinishFunc with different
+// avifEncoderCustomEncodeImageItem are not thread-safe.
+typedef avifResult (*avifEncoderCustomEncodeImageFunc)(struct avifEncoder * encoder,
+                                                       const avifImage * image,
+                                                       const avifEncoderCustomEncodeImageItem * item,
+                                                       const struct avifEncoderCustomEncodeImageArgs * args);
+
+// Called for each coded image item (color, alpha, or gain map) if avifEncoderCustomEncodeImageFunc
+// returned AVIF_RESULT_OK for the same avifEncoderCustomEncodeImageItem.
+// Called in a loop as long as it returns 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,
+                                                        const avifEncoderCustomEncodeImageItem * item,
+                                                        avifROData * sample);
+
 // ---------------------------------------------------------------------------
 // avifEncoder
 
@@ -1495,6 +1559,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 (and track if any).
+    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
@@ -1522,7 +1593,6 @@ typedef enum avifAddImageFlag
     // This is enabled automatically when using the avifEncoderWrite() single-image encode path.
     AVIF_ADD_IMAGE_FLAG_SINGLE = (1 << 1)
 } avifAddImageFlag;
-typedef uint32_t avifAddImageFlags;
 
 // Multi-function alternative to avifEncoderWrite() for advanced features.
 //

include/avif/internal.h

@@ -380,6 +380,7 @@ avifResult avifImageScaleWithLimit(avifImage * image,
 // ---------------------------------------------------------------------------
 // AVIF item category
 
+// TODO(yguyon): Reuse avifEncoderCustomEncodeImageItemType instead?
 typedef enum avifItemCategory
 {
     AVIF_ITEM_COLOR,

src/write.c

@@ -186,6 +186,15 @@ typedef struct avifEncoderItem
 } avifEncoderItem;
 AVIF_ARRAY_DECLARE(avifEncoderItemArray, avifEncoderItem, item);
 
+avifEncoderCustomEncodeImageItem avifEncoderCustomEncodeImageItemFrom(const avifEncoderItem * item)
+{
+    avifEncoderCustomEncodeImageItem value = { (avifEncoderCustomEncodeImageItemType)item->itemCategory,
+                                               item->id,
+                                               item->gridCols ? item->cellIndex / item->gridCols : 0,
+                                               item->gridCols ? item->cellIndex % item->gridCols : 0 };
+    return value;
+}
+
 // ---------------------------------------------------------------------------
 // avifEncoderItemReference
 
@@ -243,6 +252,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 +2115,35 @@ 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 avifEncoderCustomEncodeImageItem current_item = avifEncoderCustomEncodeImageItemFrom(item);
+                const avifEncoderCustomEncodeImageArgs args = {
+                    addImageFlags,
+                    quantizer,
+                    encoder->data->tileRowsLog2,
+                    encoder->data->tileColsLog2,
+                };
+                encodeResult = encoder->customEncodeImageFunc(encoder, cellImage, &current_item, &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 +3123,15 @@ 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) {
+                const avifEncoderCustomEncodeImageItem current_item = avifEncoderCustomEncodeImageItemFrom(item);
+                avifROData sample = AVIF_DATA_EMPTY;
+                avifResult encodeResult;
+                while ((encodeResult = encoder->customEncodeFinishFunc(encoder, &current_item, &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,85 @@
+// 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 avifEncoderCustomEncodeImageItem* item,
+                                 const avifEncoderCustomEncodeImageArgs*) {
+  if (item->type != AVIF_ENCODER_CUSTOM_ENCODE_ITEM_COLOR ||
+      item->gridRow != 0 || item->gridColumn != 0) {
+    // Unexpected item.
+    return AVIF_RESULT_INTERNAL_ERROR;
+  }
+
+  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,
+                                  const avifEncoderCustomEncodeImageItem* item,
+                                  avifROData* sample) {
+  if (item->type != AVIF_ENCODER_CUSTOM_ENCODE_ITEM_COLOR ||
+      item->gridRow != 0 || item->gridColumn != 0) {
+    // Unexpected item.
+    return AVIF_RESULT_INTERNAL_ERROR;
+  }
+
+  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