aom: rtc: Fix to external resize with use_psnr

From f8c91a2080beb6a98c909435fb5937f82a8ef83f Mon Sep 17 00:00:00 2001
From: Marco Paniconi <[EMAIL REDACTED]>
Date: Wed, 24 Jan 2024 18:35:12 +0000
Subject: [PATCH] rtc: Fix to external resize with use_psnr

Fix an issue with external resize with use_psnr,
which is caused by the use_rtc_tf feature.

Fix is to use crop_width/height for the aom_yv12_copy
of source into orig_source, as the psnr_calculation
uses crop_width/height.

Also condition the realloc on change in current
resolution vs last encoded one, and remove the
check on key frame for entering the use_rtc_tf in
encoder.c

Add unittest to trigger issue.

Bug: aomedia:3531

Change-Id: I85512bd4d91cd0aff9c94b0947c72271e0053180
---
 aom_scale/aom_scale_rtcd.pl    |  6 +++---
 aom_scale/generic/yv12extend.c | 38 ++++++++++++++++++++-------------
 av1/encoder/encode_strategy.c  |  6 +++---
 av1/encoder/encoder.c          | 22 ++++++++++---------
 av1/encoder/picklpf.c          |  6 +++---
 test/resize_test.cc            | 39 ++++++++++++++++++++++++++++++++++
 6 files changed, 83 insertions(+), 34 deletions(-)

diff --git a/aom_scale/aom_scale_rtcd.pl b/aom_scale/aom_scale_rtcd.pl
index ae0a85687f..0e65f54ce1 100644
--- a/aom_scale/aom_scale_rtcd.pl
+++ b/aom_scale/aom_scale_rtcd.pl
@@ -32,11 +32,11 @@ ()
 
 add_proto qw/void aom_yv12_copy_frame/, "const struct yv12_buffer_config *src_bc, struct yv12_buffer_config *dst_bc, const int num_planes";
 
-add_proto qw/void aom_yv12_copy_y/, "const struct yv12_buffer_config *src_ybc, struct yv12_buffer_config *dst_ybc";
+add_proto qw/void aom_yv12_copy_y/, "const struct yv12_buffer_config *src_ybc, struct yv12_buffer_config *dst_ybc, int use_crop";
 
-add_proto qw/void aom_yv12_copy_u/, "const struct yv12_buffer_config *src_bc, struct yv12_buffer_config *dst_bc";
+add_proto qw/void aom_yv12_copy_u/, "const struct yv12_buffer_config *src_bc, struct yv12_buffer_config *dst_bc, int use_crop";
 
-add_proto qw/void aom_yv12_copy_v/, "const struct yv12_buffer_config *src_bc, struct yv12_buffer_config *dst_bc";
+add_proto qw/void aom_yv12_copy_v/, "const struct yv12_buffer_config *src_bc, struct yv12_buffer_config *dst_bc, int use_crop";
 
 add_proto qw/void aom_yv12_partial_copy_y/, "const struct yv12_buffer_config *src_ybc, int hstart1, int hend1, int vstart1, int vend1, struct yv12_buffer_config *dst_ybc, int hstart2, int vstart2";
 add_proto qw/void aom_yv12_partial_coloc_copy_y/, "const struct yv12_buffer_config *src_ybc, struct yv12_buffer_config *dst_ybc, int hstart, int hend, int vstart, int vend";
diff --git a/aom_scale/generic/yv12extend.c b/aom_scale/generic/yv12extend.c
index 5546112d40..61e492cf06 100644
--- a/aom_scale/generic/yv12extend.c
+++ b/aom_scale/generic/yv12extend.c
@@ -301,9 +301,13 @@ void aom_yv12_copy_frame_c(const YV12_BUFFER_CONFIG *src_bc,
   aom_yv12_extend_frame_borders_c(dst_bc, num_planes);
 }
 
+// TODO(marpan/wtc): Look into why use_crop can't always be 1.
+// Some tests are currently failing if 1 is used everywhere.
 void aom_yv12_copy_y_c(const YV12_BUFFER_CONFIG *src_ybc,
-                       YV12_BUFFER_CONFIG *dst_ybc) {
+                       YV12_BUFFER_CONFIG *dst_ybc, int use_crop) {
   int row;
+  int width = use_crop ? src_ybc->y_crop_width : src_ybc->y_width;
+  int height = use_crop ? src_ybc->y_crop_height : src_ybc->y_height;
   const uint8_t *src = src_ybc->y_buffer;
   uint8_t *dst = dst_ybc->y_buffer;
 
@@ -311,8 +315,8 @@ void aom_yv12_copy_y_c(const YV12_BUFFER_CONFIG *src_ybc,
   if (src_ybc->flags & YV12_FLAG_HIGHBITDEPTH) {
     const uint16_t *src16 = CONVERT_TO_SHORTPTR(src);
     uint16_t *dst16 = CONVERT_TO_SHORTPTR(dst);
-    for (row = 0; row < src_ybc->y_height; ++row) {
-      memcpy(dst16, src16, src_ybc->y_width * sizeof(uint16_t));
+    for (row = 0; row < height; ++row) {
+      memcpy(dst16, src16, width * sizeof(uint16_t));
       src16 += src_ybc->y_stride;
       dst16 += dst_ybc->y_stride;
     }
@@ -320,56 +324,60 @@ void aom_yv12_copy_y_c(const YV12_BUFFER_CONFIG *src_ybc,
   }
 #endif
 
-  for (row = 0; row < src_ybc->y_height; ++row) {
-    memcpy(dst, src, src_ybc->y_width);
+  for (row = 0; row < height; ++row) {
+    memcpy(dst, src, width);
     src += src_ybc->y_stride;
     dst += dst_ybc->y_stride;
   }
 }
 
 void aom_yv12_copy_u_c(const YV12_BUFFER_CONFIG *src_bc,
-                       YV12_BUFFER_CONFIG *dst_bc) {
+                       YV12_BUFFER_CONFIG *dst_bc, int use_crop) {
   int row;
+  int width = use_crop ? src_bc->uv_crop_width : src_bc->uv_width;
+  int height = use_crop ? src_bc->uv_crop_height : src_bc->uv_height;
   const uint8_t *src = src_bc->u_buffer;
   uint8_t *dst = dst_bc->u_buffer;
 #if CONFIG_AV1_HIGHBITDEPTH
   if (src_bc->flags & YV12_FLAG_HIGHBITDEPTH) {
     const uint16_t *src16 = CONVERT_TO_SHORTPTR(src);
     uint16_t *dst16 = CONVERT_TO_SHORTPTR(dst);
-    for (row = 0; row < src_bc->uv_height; ++row) {
-      memcpy(dst16, src16, src_bc->uv_width * sizeof(uint16_t));
+    for (row = 0; row < height; ++row) {
+      memcpy(dst16, src16, width * sizeof(uint16_t));
       src16 += src_bc->uv_stride;
       dst16 += dst_bc->uv_stride;
     }
     return;
   }
 #endif
-  for (row = 0; row < src_bc->uv_height; ++row) {
-    memcpy(dst, src, src_bc->uv_width);
+  for (row = 0; row < height; ++row) {
+    memcpy(dst, src, width);
     src += src_bc->uv_stride;
     dst += dst_bc->uv_stride;
   }
 }
 
 void aom_yv12_copy_v_c(const YV12_BUFFER_CONFIG *src_bc,
-                       YV12_BUFFER_CONFIG *dst_bc) {
+                       YV12_BUFFER_CONFIG *dst_bc, int use_crop) {
   int row;
+  int width = use_crop ? src_bc->uv_crop_width : src_bc->uv_width;
+  int height = use_crop ? src_bc->uv_crop_height : src_bc->uv_height;
   const uint8_t *src = src_bc->v_buffer;
   uint8_t *dst = dst_bc->v_buffer;
 #if CONFIG_AV1_HIGHBITDEPTH
   if (src_bc->flags & YV12_FLAG_HIGHBITDEPTH) {
     const uint16_t *src16 = CONVERT_TO_SHORTPTR(src);
     uint16_t *dst16 = CONVERT_TO_SHORTPTR(dst);
-    for (row = 0; row < src_bc->uv_height; ++row) {
-      memcpy(dst16, src16, src_bc->uv_width * sizeof(uint16_t));
+    for (row = 0; row < height; ++row) {
+      memcpy(dst16, src16, width * sizeof(uint16_t));
       src16 += src_bc->uv_stride;
       dst16 += dst_bc->uv_stride;
     }
     return;
   }
 #endif
-  for (row = 0; row < src_bc->uv_height; ++row) {
-    memcpy(dst, src, src_bc->uv_width);
+  for (row = 0; row < height; ++row) {
+    memcpy(dst, src, width);
     src += src_bc->uv_stride;
     dst += dst_bc->uv_stride;
   }
diff --git a/av1/encoder/encode_strategy.c b/av1/encoder/encode_strategy.c
index 35ca83c3f4..07832a94b7 100644
--- a/av1/encoder/encode_strategy.c
+++ b/av1/encoder/encode_strategy.c
@@ -1758,9 +1758,9 @@ int av1_encode_strategy(AV1_COMP *const cpi, size_t *const size,
       cpi->svc.temporal_layer_id == 0 &&
       cpi->unscaled_source->y_width == cpi->svc.source_last_TL0.y_width &&
       cpi->unscaled_source->y_height == cpi->svc.source_last_TL0.y_height) {
-    aom_yv12_copy_y(cpi->unscaled_source, &cpi->svc.source_last_TL0);
-    aom_yv12_copy_u(cpi->unscaled_source, &cpi->svc.source_last_TL0);
-    aom_yv12_copy_v(cpi->unscaled_source, &cpi->svc.source_last_TL0);
+    aom_yv12_copy_y(cpi->unscaled_source, &cpi->svc.source_last_TL0, 0);
+    aom_yv12_copy_u(cpi->unscaled_source, &cpi->svc.source_last_TL0, 0);
+    aom_yv12_copy_v(cpi->unscaled_source, &cpi->svc.source_last_TL0, 0);
   }
 
   return AOM_CODEC_OK;
diff --git a/av1/encoder/encoder.c b/av1/encoder/encoder.c
index fe053af5cc..fc33675af7 100644
--- a/av1/encoder/encoder.c
+++ b/av1/encoder/encoder.c
@@ -2667,13 +2667,12 @@ static int encode_without_recode(AV1_COMP *cpi) {
   cm->cur_frame->seg.enabled = cm->seg.enabled;
 
   // This is for rtc temporal filtering case.
-  if (is_psnr_calc_enabled(cpi) && cpi->sf.rt_sf.use_rtc_tf &&
-      cm->current_frame.frame_type != KEY_FRAME) {
+  if (is_psnr_calc_enabled(cpi) && cpi->sf.rt_sf.use_rtc_tf) {
     const SequenceHeader *seq_params = cm->seq_params;
 
     if (cpi->orig_source.buffer_alloc_sz == 0 ||
-        cpi->last_source->y_width != cpi->source->y_width ||
-        cpi->last_source->y_height != cpi->source->y_height) {
+        cpi->rc.prev_coded_width != cpi->oxcf.frm_dim_cfg.width ||
+        cpi->rc.prev_coded_height != cpi->oxcf.frm_dim_cfg.height) {
       // Allocate a source buffer to store the true source for psnr calculation.
       if (aom_alloc_frame_buffer(
               &cpi->orig_source, cpi->oxcf.frm_dim_cfg.width,
@@ -2684,9 +2683,12 @@ static int encode_without_recode(AV1_COMP *cpi) {
                            "Failed to allocate scaled buffer");
     }
 
-    aom_yv12_copy_y(cpi->source, &cpi->orig_source);
-    aom_yv12_copy_u(cpi->source, &cpi->orig_source);
-    aom_yv12_copy_v(cpi->source, &cpi->orig_source);
+    // Copy source into orig_source, needed for the psnr computation.
+    // Use crop_width/height for aom_yv12_copy, as the psnr computation
+    // uses crop_width/height.
+    aom_yv12_copy_y(cpi->source, &cpi->orig_source, 1);
+    aom_yv12_copy_u(cpi->source, &cpi->orig_source, 1);
+    aom_yv12_copy_v(cpi->source, &cpi->orig_source, 1);
   }
 
 #if CONFIG_COLLECT_COMPONENT_TIMING
@@ -2725,9 +2727,9 @@ static int encode_without_recode(AV1_COMP *cpi) {
       (cm->width != cpi->unscaled_source->y_crop_width ||
        cm->height != cpi->unscaled_source->y_crop_height)) {
     cpi->scaled_last_source_available = 1;
-    aom_yv12_copy_y(&cpi->scaled_source, &cpi->scaled_last_source);
-    aom_yv12_copy_u(&cpi->scaled_source, &cpi->scaled_last_source);
-    aom_yv12_copy_v(&cpi->scaled_source, &cpi->scaled_last_source);
+    aom_yv12_copy_y(&cpi->scaled_source, &cpi->scaled_last_source, 0);
+    aom_yv12_copy_u(&cpi->scaled_source, &cpi->scaled_last_source, 0);
+    aom_yv12_copy_v(&cpi->scaled_source, &cpi->scaled_last_source, 0);
   }
 
 #if CONFIG_COLLECT_COMPONENT_TIMING
diff --git a/av1/encoder/picklpf.c b/av1/encoder/picklpf.c
index 9084d3f13a..9963cb07d0 100644
--- a/av1/encoder/picklpf.c
+++ b/av1/encoder/picklpf.c
@@ -30,9 +30,9 @@
 static void yv12_copy_plane(const YV12_BUFFER_CONFIG *src_bc,
                             YV12_BUFFER_CONFIG *dst_bc, int plane) {
   switch (plane) {
-    case 0: aom_yv12_copy_y(src_bc, dst_bc); break;
-    case 1: aom_yv12_copy_u(src_bc, dst_bc); break;
-    case 2: aom_yv12_copy_v(src_bc, dst_bc); break;
+    case 0: aom_yv12_copy_y(src_bc, dst_bc, 0); break;
+    case 1: aom_yv12_copy_u(src_bc, dst_bc, 0); break;
+    case 2: aom_yv12_copy_v(src_bc, dst_bc, 0); break;
     default: assert(plane >= 0 && plane <= 2); break;
   }
 }
diff --git a/test/resize_test.cc b/test/resize_test.cc
index dd66f83557..a84a4654a8 100644
--- a/test/resize_test.cc
+++ b/test/resize_test.cc
@@ -689,6 +689,45 @@ TEST_P(ResizeRealtimeTest, TestExternalResizeWorks) {
   }
 }
 
+TEST_P(ResizeRealtimeTest, TestExternalResizeWorksUsePSNR) {
+  ResizingVideoSource video;
+  video.flag_codec_ = 1;
+  change_bitrate_ = false;
+  set_scale_mode_ = false;
+  set_scale_mode2_ = false;
+  set_scale_mode3_ = false;
+  mismatch_psnr_ = 0.0;
+  mismatch_nframes_ = 0;
+  init_flags_ = AOM_CODEC_USE_PSNR;
+  cfg_.rc_dropframe_thresh = 30;
+  DefaultConfig();
+  // Test external resizing with start resolution equal to
+  // 1. kInitialWidth and kInitialHeight
+  // 2. down-scaled kInitialWidth and kInitialHeight
+  for (int i = 0; i < 2; i++) {
+    video.change_start_resln_ = static_cast<bool>(i);
+
+    ASSERT_NO_FATAL_FAILURE(RunLoop(&video));
+
+    // Check we decoded the same number of frames as we attempted to encode
+    ASSERT_EQ(frame_info_list_.size(), video.limit());
+    for (const auto &info : frame_info_list_) {
+      const unsigned int frame = static_cast<unsigned>(info.pts);
+      unsigned int expected_w;
+      unsigned int expected_h;
+      ScaleForFrameNumber(frame, kInitialWidth, kInitialHeight,
+                          video.flag_codec_, video.change_start_resln_,
+                          &expected_w, &expected_h);
+      EXPECT_EQ(expected_w, info.w)
+          << "Frame " << frame << " had unexpected width";
+      EXPECT_EQ(expected_h, info.h)
+          << "Frame " << frame << " had unexpected height";
+      EXPECT_EQ(static_cast<unsigned int>(0), GetMismatchFrames());
+    }
+    frame_info_list_.clear();
+  }
+}
+
 // Verify the dynamic resizer behavior for real time, 1 pass CBR mode.
 // Run at low bitrate, with resize_allowed = 1, and verify that we get
 // one resize down event.