Add sato BIT_DEPTH_EXTENSION_12B_8B_OVERLAP_4B (#2215)

y-guyon · web-flow · commit c898b5eb9873 · 2024-08-05T14:03:38.000+02:00
Sample Transform derived image item support for residual bit depth
extension.
diff --git a/include/avif/avif.h b/include/avif/avif.h
@@ -746,7 +746,16 @@ typedef enum avifSampleTransformRecipe
     // (ignoring the hidden image item), leading to a valid image but with
     // precision loss (16-bit samples truncated to the 12 most significant
     // bits).
-    AVIF_SAMPLE_TRANSFORM_BIT_DEPTH_EXTENSION_12B_4B
+    AVIF_SAMPLE_TRANSFORM_BIT_DEPTH_EXTENSION_12B_4B,
+    // Encode the 12 most significant bits of each input image sample lossily or
+    // losslessly into a base image. The difference between the original and
+    // decoded values of these samples is encoded as a separate 8-bit hidden
+    // image item. The two are combined at decoding into one image with the same
+    // bit depth as the original image. It is backward compatible in the sense
+    // that it is possible to decode only the base image (ignoring the hidden
+    // image item), leading to a valid image but with loss due to precision
+    // truncation and/or compression.
+    AVIF_SAMPLE_TRANSFORM_BIT_DEPTH_EXTENSION_12B_8B_OVERLAP_4B
 } avifSampleTransformRecipe;
 #endif // AVIF_ENABLE_EXPERIMENTAL_SAMPLE_TRANSFORM
 
diff --git a/src/sampletransform.c b/src/sampletransform.c
@@ -137,14 +137,42 @@ avifResult avifSampleTransformRecipeToExpression(avifSampleTransformRecipe recip
         return AVIF_RESULT_OK;
     }
 
+    if (recipe == AVIF_SAMPLE_TRANSFORM_BIT_DEPTH_EXTENSION_12B_8B_OVERLAP_4B) {
+        // reference_count is two: one 12-bit input image and one 8-bit input image.
+        //   (base_sample << 4) + hidden_sample
+        // Note: Both base_sample and hidden_sample are encoded lossily or losslessly. hidden_sample overlaps
+        //       with base_sample by 4 bits to alleviate the loss caused by the quantization of base_sample.
+        AVIF_CHECKERR(avifArrayCreate(expression, sizeof(avifSampleTransformToken), 7), AVIF_RESULT_OUT_OF_MEMORY);
+
+        // The base image represents the 12 most significant bits of the reconstructed, bit-depth-extended output image.
+        // Left shift the base image (which is also the primary item, or the auxiliary alpha item of the primary item)
+        // by 4 bits. This is equivalent to multiplying by 2^4.
+        AVIF_ASSERT_OR_RETURN(avifPushConstant(expression, 16));
+        AVIF_ASSERT_OR_RETURN(avifPushInputImageItem(expression, 1));
+        AVIF_ASSERT_OR_RETURN(avifPushOperator(expression, AVIF_SAMPLE_TRANSFORM_PRODUCT));
+
+        // The second image represents the offset to apply to the shifted base image to retrieve
+        // the original image, with some loss due to quantization.
+        AVIF_ASSERT_OR_RETURN(avifPushInputImageItem(expression, 2));
+        AVIF_ASSERT_OR_RETURN(avifPushOperator(expression, AVIF_SAMPLE_TRANSFORM_SUM));
+
+        // The second image is offset by 128 to have unsigned values to encode.
+        // Correct that last to always work with unsigned values in the operations above.
+        AVIF_ASSERT_OR_RETURN(avifPushConstant(expression, 128));
+        AVIF_ASSERT_OR_RETURN(avifPushOperator(expression, AVIF_SAMPLE_TRANSFORM_DIFFERENCE));
+        // Sample values are clamped to [0:1<<depth[ at that point.
+        return AVIF_RESULT_OK;
+    }
+
     return AVIF_RESULT_INVALID_ARGUMENT;
 }
 
 avifResult avifSampleTransformExpressionToRecipe(const avifSampleTransformExpression * expression, avifSampleTransformRecipe * recipe)
 {
     *recipe = AVIF_SAMPLE_TRANSFORM_NONE;
     const avifSampleTransformRecipe kAllRecipes[] = { AVIF_SAMPLE_TRANSFORM_BIT_DEPTH_EXTENSION_8B_8B,
-                                                      AVIF_SAMPLE_TRANSFORM_BIT_DEPTH_EXTENSION_12B_4B };
+                                                      AVIF_SAMPLE_TRANSFORM_BIT_DEPTH_EXTENSION_12B_4B,
+                                                      AVIF_SAMPLE_TRANSFORM_BIT_DEPTH_EXTENSION_12B_8B_OVERLAP_4B };
     for (size_t i = 0; i < sizeof(kAllRecipes) / sizeof(kAllRecipes[0]); ++i) {
         avifSampleTransformRecipe candidateRecipe = kAllRecipes[i];
         avifSampleTransformExpression candidateExpression = { 0 };
diff --git a/src/write.c b/src/write.c
@@ -1210,7 +1210,8 @@ static avifResult avifEncoderCreateBitDepthExtensionItems(avifEncoder * encoder,
                                                           uint16_t colorItemID)
 {
     AVIF_ASSERT_OR_RETURN(encoder->sampleTransformRecipe == AVIF_SAMPLE_TRANSFORM_BIT_DEPTH_EXTENSION_8B_8B ||
-                          encoder->sampleTransformRecipe == AVIF_SAMPLE_TRANSFORM_BIT_DEPTH_EXTENSION_12B_4B);
+                          encoder->sampleTransformRecipe == AVIF_SAMPLE_TRANSFORM_BIT_DEPTH_EXTENSION_12B_4B ||
+                          encoder->sampleTransformRecipe == AVIF_SAMPLE_TRANSFORM_BIT_DEPTH_EXTENSION_12B_8B_OVERLAP_4B);
 
     // There are multiple possible ISOBMFF box hierarchies for translucent images,
     // using 'sato' (Sample Transform) derived image items:
@@ -1300,7 +1301,52 @@ static avifResult avifImageCreateAllocate(avifImage ** sampleTransformedImage, c
     return avifImageAllocatePlanes(*sampleTransformedImage, planes);
 }
 
-static avifResult avifEncoderCreateSatoImage(const avifEncoder * encoder,
+// Finds the encoded base image and decodes it. Callers of this function must free
+// *codec and *decodedBaseImage if not null, whether the function succeeds or not.
+static avifResult avifEncoderDecodeSatoBaseImage(avifEncoder * encoder,
+                                                 const avifImage * original,
+                                                 uint32_t numBits,
+                                                 avifPlanesFlag planes,
+                                                 avifCodec ** codec,
+                                                 avifImage ** decodedBaseImage)
+{
+    avifDecodeSample sample;
+    memset(&sample, 0, sizeof(sample));
+    sample.spatialID = AVIF_SPATIAL_ID_UNSET;
+
+    // Find the encoded bytes of the base image item.
+    for (uint32_t itemIndex = 0; itemIndex < encoder->data->items.count; ++itemIndex) {
+        avifEncoderItem * item = &encoder->data->items.item[itemIndex];
+        if ((item->itemCategory != AVIF_ITEM_COLOR || planes != AVIF_PLANES_YUV) &&
+            (item->itemCategory != AVIF_ITEM_ALPHA || planes != AVIF_PLANES_A)) {
+            continue;
+        }
+
+        AVIF_ASSERT_OR_RETURN(item->encodeOutput != NULL); // TODO: Support grids?
+        AVIF_ASSERT_OR_RETURN(item->encodeOutput->samples.count == 1);
+        AVIF_ASSERT_OR_RETURN(item->encodeOutput->samples.sample[0].data.size != 0);
+        AVIF_ASSERT_OR_RETURN(sample.data.size == 0); // There should be only one base item.
+        sample.data.data = item->encodeOutput->samples.sample[0].data.data;
+        sample.data.size = item->encodeOutput->samples.sample[0].data.size;
+    }
+    AVIF_ASSERT_OR_RETURN(sample.data.size != 0); // There should be at least one base item.
+
+    // avifCodecGetNextImageFunc() uses only a few fields of its decoder argument.
+    avifDecoder decoder;
+    memset(&decoder, 0, sizeof(decoder));
+    decoder.maxThreads = encoder->maxThreads;
+    decoder.imageSizeLimit = AVIF_DEFAULT_IMAGE_SIZE_LIMIT;
+
+    AVIF_CHECKRES(avifCodecCreate(AVIF_CODEC_CHOICE_AUTO, AVIF_CODEC_FLAG_CAN_DECODE, codec));
+    (*codec)->diag = &encoder->diag;
+    AVIF_CHECKRES(avifImageCreateAllocate(decodedBaseImage, original, numBits, planes));
+    avifBool isLimitedRangeAlpha = AVIF_FALSE; // Ignored.
+    AVIF_CHECKERR((*codec)->getNextImage(*codec, &decoder, &sample, planes == AVIF_PLANES_A, &isLimitedRangeAlpha, *decodedBaseImage),
+                  AVIF_RESULT_ENCODE_SAMPLE_TRANSFORM_FAILED);
+    return AVIF_RESULT_OK;
+}
+
+static avifResult avifEncoderCreateSatoImage(avifEncoder * encoder,
                                              const avifEncoderItem * item,
                                              avifBool itemWillBeEncodedLosslessly,
                                              const avifImage * image,
@@ -1323,8 +1369,7 @@ static avifResult avifEncoderCreateSatoImage(const avifEncoder * encoder,
             AVIF_CHECKRES(avifImageCreateAllocate(sampleTransformedImage, image, 8, planes));
             AVIF_CHECKRES(avifImageApplyImgOpConst(*sampleTransformedImage, image, AVIF_SAMPLE_TRANSFORM_AND, 255, planes));
         }
-    } else {
-        AVIF_CHECKERR(encoder->sampleTransformRecipe == AVIF_SAMPLE_TRANSFORM_BIT_DEPTH_EXTENSION_12B_4B, AVIF_RESULT_NOT_IMPLEMENTED);
+    } else if (encoder->sampleTransformRecipe == AVIF_SAMPLE_TRANSFORM_BIT_DEPTH_EXTENSION_12B_4B) {
         if (isBase) {
             AVIF_CHECKRES(avifImageCreateAllocate(sampleTransformedImage, image, 12, planes));
             AVIF_CHECKRES(avifImageApplyImgOpConst(*sampleTransformedImage, image, AVIF_SAMPLE_TRANSFORM_DIVIDE, 16, planes));
@@ -1351,11 +1396,49 @@ static avifResult avifEncoderCreateSatoImage(const avifEncoder * encoder,
                     avifImageApplyImgOpConst(*sampleTransformedImage, *sampleTransformedImage, AVIF_SAMPLE_TRANSFORM_SUM, 7, planes));
             }
         }
+    } else {
+        AVIF_CHECKERR(encoder->sampleTransformRecipe == AVIF_SAMPLE_TRANSFORM_BIT_DEPTH_EXTENSION_12B_8B_OVERLAP_4B,
+                      AVIF_RESULT_NOT_IMPLEMENTED);
+        if (isBase) {
+            AVIF_CHECKRES(avifImageCreateAllocate(sampleTransformedImage, image, 12, planes));
+            AVIF_CHECKRES(avifImageApplyImgOpConst(*sampleTransformedImage, image, AVIF_SAMPLE_TRANSFORM_DIVIDE, 16, planes));
+        } else {
+            AVIF_CHECKRES(avifImageCreateAllocate(sampleTransformedImage, image, 8, planes));
+            avifCodec * codec = NULL;
+            avifImage * decodedBaseImage = NULL;
+            avifResult result = avifEncoderDecodeSatoBaseImage(encoder, image, 12, planes, &codec, &decodedBaseImage);
+            if (result == AVIF_RESULT_OK) {
+                // decoded = main*16+hidden-128 so hidden = clamp_8b(original-main*16+128). Postfix notation.
+                const avifSampleTransformToken tokens[] = { { AVIF_SAMPLE_TRANSFORM_INPUT_IMAGE_ITEM_INDEX, 0, /*inputImageItemIndex=*/1 },
+                                                            { AVIF_SAMPLE_TRANSFORM_INPUT_IMAGE_ITEM_INDEX, 0, /*inputImageItemIndex=*/2 },
+                                                            { AVIF_SAMPLE_TRANSFORM_CONSTANT, /*constant=*/16, 0 },
+                                                            { AVIF_SAMPLE_TRANSFORM_PRODUCT, 0, 0 },
+                                                            { AVIF_SAMPLE_TRANSFORM_DIFFERENCE, 0, 0 },
+                                                            { AVIF_SAMPLE_TRANSFORM_CONSTANT, /*constant=*/128, 0 },
+                                                            { AVIF_SAMPLE_TRANSFORM_SUM, 0, 0 } };
+                // image is "original" (index 1) and decodedBaseImage is "main" (index 2) in the formula above.
+                const avifImage * inputImageItems[] = { image, decodedBaseImage };
+                result = avifImageApplyOperations(*sampleTransformedImage,
+                                                  AVIF_SAMPLE_TRANSFORM_BIT_DEPTH_32,
+                                                  /*numTokens=*/7,
+                                                  tokens,
+                                                  /*numInputImageItems=*/2,
+                                                  inputImageItems,
+                                                  planes);
+            }
+            if (decodedBaseImage) {
+                avifImageDestroy(decodedBaseImage);
+            }
+            if (codec) {
+                avifCodecDestroy(codec);
+            }
+            AVIF_CHECKRES(result);
+        }
     }
     return AVIF_RESULT_OK;
 }
 
-static avifResult avifEncoderCreateBitDepthExtensionImage(const avifEncoder * encoder,
+static avifResult avifEncoderCreateBitDepthExtensionImage(avifEncoder * encoder,
                                                           const avifEncoderItem * item,
                                                           avifBool itemWillBeEncodedLosslessly,
                                                           const avifImage * image,
@@ -1807,7 +1890,8 @@ static avifResult avifEncoderAddImageInternal(avifEncoder * encoder,
 
 #if defined(AVIF_ENABLE_EXPERIMENTAL_SAMPLE_TRANSFORM)
         if (encoder->sampleTransformRecipe == AVIF_SAMPLE_TRANSFORM_BIT_DEPTH_EXTENSION_8B_8B ||
-            encoder->sampleTransformRecipe == AVIF_SAMPLE_TRANSFORM_BIT_DEPTH_EXTENSION_12B_4B) {
+            encoder->sampleTransformRecipe == AVIF_SAMPLE_TRANSFORM_BIT_DEPTH_EXTENSION_12B_4B ||
+            encoder->sampleTransformRecipe == AVIF_SAMPLE_TRANSFORM_BIT_DEPTH_EXTENSION_12B_8B_OVERLAP_4B) {
             // For now, only 16-bit depth is supported.
             AVIF_ASSERT_OR_RETURN(firstCell->depth == 16);
 #if defined(AVIF_ENABLE_EXPERIMENTAL_GAIN_MAP)
@@ -2655,7 +2739,8 @@ static avifResult avifRWStreamWriteProperties(avifItemPropertyDedup * const dedu
         uint8_t depth = (uint8_t)itemMetadata->depth;
 #if defined(AVIF_ENABLE_EXPERIMENTAL_SAMPLE_TRANSFORM)
         if (encoder->sampleTransformRecipe == AVIF_SAMPLE_TRANSFORM_BIT_DEPTH_EXTENSION_8B_8B ||
-            encoder->sampleTransformRecipe == AVIF_SAMPLE_TRANSFORM_BIT_DEPTH_EXTENSION_12B_4B) {
+            encoder->sampleTransformRecipe == AVIF_SAMPLE_TRANSFORM_BIT_DEPTH_EXTENSION_12B_4B ||
+            encoder->sampleTransformRecipe == AVIF_SAMPLE_TRANSFORM_BIT_DEPTH_EXTENSION_12B_8B_OVERLAP_4B) {
             if (item->itemCategory == AVIF_ITEM_SAMPLE_TRANSFORM) {
                 AVIF_ASSERT_OR_RETURN(depth == 16); // Only 16-bit depth is supported for now.
             } else if (encoder->sampleTransformRecipe == AVIF_SAMPLE_TRANSFORM_BIT_DEPTH_EXTENSION_8B_8B) {
diff --git a/tests/gtest/avif16bittest.cc b/tests/gtest/avif16bittest.cc
@@ -76,21 +76,25 @@ TEST_P(SampleTransformTest, Avif16bit) {
       image->alphaPlane ? AVIF_PLANES_ALL : AVIF_PLANES_YUV, image->yuvRange);
   ASSERT_NE(image_no_sato, nullptr);
 
-  const uint32_t shift =
-      recipe == AVIF_SAMPLE_TRANSFORM_BIT_DEPTH_EXTENSION_8B_8B ? 8 : 4;
-  const avifImage* inputImage = image.get();
-  // Postfix notation.
-  const avifSampleTransformToken tokens[] = {
-      {AVIF_SAMPLE_TRANSFORM_INPUT_IMAGE_ITEM_INDEX, 0,
-       /*inputImageItemIndex=*/1},
-      {AVIF_SAMPLE_TRANSFORM_CONSTANT, 1 << shift, 0},
-      {AVIF_SAMPLE_TRANSFORM_DIVIDE, 0, 0}};
-  ASSERT_EQ(avifImageApplyOperations(
-                image_no_sato.get(), AVIF_SAMPLE_TRANSFORM_BIT_DEPTH_32,
-                /*numTokens=*/3, tokens, /*numInputImageItems=*/1, &inputImage,
-                AVIF_PLANES_ALL),
-            AVIF_RESULT_OK);
-  EXPECT_TRUE(testutil::AreImagesEqual(*image_no_sato, *decoded_no_sato));
+  if (recipe == AVIF_SAMPLE_TRANSFORM_BIT_DEPTH_EXTENSION_8B_8B ||
+      recipe == AVIF_SAMPLE_TRANSFORM_BIT_DEPTH_EXTENSION_12B_4B) {
+    // These recipes always encode the primary item losslessly. Check that.
+    const uint32_t shift =
+        recipe == AVIF_SAMPLE_TRANSFORM_BIT_DEPTH_EXTENSION_8B_8B ? 8 : 4;
+    const avifImage* inputImage = image.get();
+    // Postfix notation.
+    const avifSampleTransformToken tokens[] = {
+        {AVIF_SAMPLE_TRANSFORM_INPUT_IMAGE_ITEM_INDEX, 0,
+         /*inputImageItemIndex=*/1},
+        {AVIF_SAMPLE_TRANSFORM_CONSTANT, 1 << shift, 0},
+        {AVIF_SAMPLE_TRANSFORM_DIVIDE, 0, 0}};
+    ASSERT_EQ(avifImageApplyOperations(
+                  image_no_sato.get(), AVIF_SAMPLE_TRANSFORM_BIT_DEPTH_32,
+                  /*numTokens=*/3, tokens, /*numInputImageItems=*/1,
+                  &inputImage, AVIF_PLANES_ALL),
+              AVIF_RESULT_OK);
+    EXPECT_TRUE(testutil::AreImagesEqual(*image_no_sato, *decoded_no_sato));
+  }
 }
 
 //------------------------------------------------------------------------------
@@ -115,6 +119,16 @@ INSTANTIATE_TEST_SUITE_P(
         /*quality=*/
         testing::Values(AVIF_QUALITY_LOSSLESS)));
 
+INSTANTIATE_TEST_SUITE_P(
+    ResidualBitDepthExtension, SampleTransformTest,
+    testing::Combine(
+        testing::Values(
+            AVIF_SAMPLE_TRANSFORM_BIT_DEPTH_EXTENSION_12B_8B_OVERLAP_4B),
+        testing::Values(AVIF_PIXEL_FORMAT_YUV444),
+        /*create_alpha=*/testing::Values(false),
+        /*quality=*/
+        testing::Values(AVIF_QUALITY_DEFAULT)));
+
 INSTANTIATE_TEST_SUITE_P(
     Alpha, SampleTransformTest,
     testing::Combine(
diff --git a/tests/gtest/avifsampletransformtest.cc b/tests/gtest/avifsampletransformtest.cc
@@ -85,7 +85,8 @@ TEST(SampleTransformTest, NoRecipe) {
 TEST(SampleTransformTest, RecipeToExpression) {
   for (avifSampleTransformRecipe recipe :
        {AVIF_SAMPLE_TRANSFORM_BIT_DEPTH_EXTENSION_8B_8B,
-        AVIF_SAMPLE_TRANSFORM_BIT_DEPTH_EXTENSION_12B_4B}) {
+        AVIF_SAMPLE_TRANSFORM_BIT_DEPTH_EXTENSION_12B_4B,
+        AVIF_SAMPLE_TRANSFORM_BIT_DEPTH_EXTENSION_12B_8B_OVERLAP_4B}) {
     AvifExpression expression;
     ASSERT_EQ(avifSampleTransformRecipeToExpression(recipe, &expression),
               AVIF_RESULT_OK);
