From 487c3d821df79178edd18a62285449d8d1f70160 Mon Sep 17 00:00:00 2001
From: Dirk Farin <dirk.farin@gmail.com>
Date: Mon, 12 Jul 2021 18:16:11 +0200
Subject: [PATCH] AVIF: signal chroma sample position when encoding (#521)

---
 libheif/heif_encoder_aom.cc   | 6 ++++++
 libheif/heif_encoder_rav1e.cc | 8 ++++++--
 2 files changed, 12 insertions(+), 2 deletions(-)

diff --git a/libheif/heif_encoder_aom.cc b/libheif/heif_encoder_aom.cc
index 65a72e29e4..6bc2f25fc2 100644
--- a/libheif/heif_encoder_aom.cc
+++ b/libheif/heif_encoder_aom.cc
@@ -590,23 +590,28 @@ struct heif_error aom_encode_image(void* encoder_raw, const struct heif_image* i
   aom_img_fmt_t img_format = AOM_IMG_FMT_NONE;
 
   int chroma_height = 0;
+  int chroma_sample_position = AOM_CSP_UNKNOWN;
 
   switch (chroma) {
     case heif_chroma_420:
     case heif_chroma_monochrome:
       img_format = AOM_IMG_FMT_I420;
       chroma_height = (source_height+1)/2;
+      chroma_sample_position = AOM_CSP_UNKNOWN; // TODO: change this to CSP_CENTER in the future (https://github.com/AOMediaCodec/av1-avif/issues/88)
       break;
     case heif_chroma_422:
       img_format = AOM_IMG_FMT_I422;
       chroma_height = (source_height+1)/2;
+      chroma_sample_position = AOM_CSP_COLOCATED;
       break;
     case heif_chroma_444:
       img_format = AOM_IMG_FMT_I444;
       chroma_height = source_height;
+      chroma_sample_position = AOM_CSP_COLOCATED;
       break;
     default:
       img_format = AOM_IMG_FMT_NONE;
+      chroma_sample_position = AOM_CSP_UNKNOWN;
       assert(false);
       break;
   }
@@ -760,6 +765,7 @@ struct heif_error aom_encode_image(void* encoder_raw, const struct heif_image* i
 
   // In aom, color_range defaults to limited range (0). Set it to full range (1).
   aom_codec_control(&codec, AV1E_SET_COLOR_RANGE, nclx ? nclx->get_full_range_flag() : 1);
+  aom_codec_control(&codec, AV1E_SET_CHROMA_SAMPLE_POSITION, chroma_sample_position);
 
   if (nclx &&
       (input_class == heif_image_input_class_normal ||
diff --git a/libheif/heif_encoder_rav1e.cc b/libheif/heif_encoder_rav1e.cc
index cec32080e4..14ce68f1dc 100644
--- a/libheif/heif_encoder_rav1e.cc
+++ b/libheif/heif_encoder_rav1e.cc
@@ -515,7 +515,9 @@ struct heif_error rav1e_encode_image(void* encoder_raw, const struct heif_image*
 
   uint8_t yShift = 0;
   RaChromaSampling chromaSampling;
+  RaChromaSamplePosition chromaPosition;
   RaPixelRange rav1eRange;
+
   if (input_class == heif_image_input_class_alpha) {
     chromaSampling = RA_CHROMA_SAMPLING_CS420; // I can't seem to get RA_CHROMA_SAMPLING_CS400 to work right now, unfortunately
   }
@@ -523,12 +525,15 @@ struct heif_error rav1e_encode_image(void* encoder_raw, const struct heif_image*
     switch (chroma) {
       case heif_chroma_444:
         chromaSampling = RA_CHROMA_SAMPLING_CS444;
+        chromaPosition = RA_CHROMA_SAMPLE_POSITION_COLOCATED;
         break;
       case heif_chroma_422:
         chromaSampling = RA_CHROMA_SAMPLING_CS422;
+        chromaPosition = RA_CHROMA_SAMPLE_POSITION_COLOCATED;
         break;
       case heif_chroma_420:
         chromaSampling = RA_CHROMA_SAMPLING_CS420;
+        chromaPosition = RA_CHROMA_SAMPLE_POSITION_UNKNOWN; // TODO: set to CENTER when AV1 and rav1e supports this
         yShift = 1;
         break;
       default:
@@ -548,8 +553,7 @@ struct heif_error rav1e_encode_image(void* encoder_raw, const struct heif_image*
   auto rav1eConfigRaw = rav1e_config_default();
   auto rav1eConfig = std::shared_ptr<RaConfig>(rav1eConfigRaw, [](RaConfig* c) { rav1e_config_unref(c); });
 
-  if (rav1e_config_set_pixel_format(rav1eConfig.get(), (uint8_t) bitDepth, chromaSampling,
-                                    RA_CHROMA_SAMPLE_POSITION_UNKNOWN, rav1eRange) < 0) {
+  if (rav1e_config_set_pixel_format(rav1eConfig.get(), (uint8_t) bitDepth, chromaSampling, chromaPosition, rav1eRange) < 0) {
     return heif_error_codec_library_error;
   }