1 files changed, 32 insertions, 12 deletions

diff --git a/third_party/aom/aom_dsp/flow_estimation/corner_detect.c b/third_party/aom/aom_dsp/flow_estimation/corner_detect.c
index 284d1bd7b8..44d423dcdf 100644
--- a/third_party/aom/aom_dsp/flow_estimation/corner_detect.c
+++ b/third_party/aom/aom_dsp/flow_estimation/corner_detect.c
@@ -20,6 +20,7 @@
 #include "aom_dsp/aom_dsp_common.h"
 #include "aom_dsp/flow_estimation/corner_detect.h"
 #include "aom_mem/aom_mem.h"
+#include "aom_util/aom_pthread.h"
 #include "av1/common/common.h"
 
 #define FAST_BARRIER 18
@@ -39,11 +40,24 @@ CornerList *av1_alloc_corner_list(void) {
   return corners;
 }
 
-static bool compute_corner_list(const ImagePyramid *pyr, CornerList *corners) {
-  const uint8_t *buf = pyr->layers[0].buffer;
-  int width = pyr->layers[0].width;
-  int height = pyr->layers[0].height;
-  int stride = pyr->layers[0].stride;
+static bool compute_corner_list(const YV12_BUFFER_CONFIG *frame, int bit_depth,
+                                int downsample_level, CornerList *corners) {
+  ImagePyramid *pyr = frame->y_pyramid;
+  const int layers =
+      aom_compute_pyramid(frame, bit_depth, downsample_level + 1, pyr);
+
+  if (layers < 0) {
+    return false;
+  }
+
+  // Clamp downsampling ratio base on max number of layers allowed
+  // for this frame size
+  downsample_level = layers - 1;
+
+  const uint8_t *buf = pyr->layers[downsample_level].buffer;
+  int width = pyr->layers[downsample_level].width;
+  int height = pyr->layers[downsample_level].height;
+  int stride = pyr->layers[downsample_level].stride;
 
   int *scores = NULL;
   int num_corners;
@@ -53,9 +67,11 @@ static bool compute_corner_list(const ImagePyramid *pyr, CornerList *corners) {
 
   if (num_corners <= MAX_CORNERS) {
     // Use all detected corners
-    if (num_corners != 0) {
-      memcpy(corners->corners, frame_corners_xy,
-             sizeof(*frame_corners_xy) * num_corners);
+    for (int i = 0; i < num_corners; i++) {
+      corners->corners[2 * i + 0] =
+          frame_corners_xy[i].x * (1 << downsample_level);
+      corners->corners[2 * i + 1] =
+          frame_corners_xy[i].y * (1 << downsample_level);
     }
     corners->num_corners = num_corners;
   } else {
@@ -85,8 +101,10 @@ static bool compute_corner_list(const ImagePyramid *pyr, CornerList *corners) {
     for (int i = 0; i < num_corners; i++) {
       if (scores[i] > threshold) {
         assert(copied_corners < MAX_CORNERS);
-        corners->corners[2 * copied_corners + 0] = frame_corners_xy[i].x;
-        corners->corners[2 * copied_corners + 1] = frame_corners_xy[i].y;
+        corners->corners[2 * copied_corners + 0] =
+            frame_corners_xy[i].x * (1 << downsample_level);
+        corners->corners[2 * copied_corners + 1] =
+            frame_corners_xy[i].y * (1 << downsample_level);
         copied_corners += 1;
       }
     }
@@ -99,7 +117,8 @@ static bool compute_corner_list(const ImagePyramid *pyr, CornerList *corners) {
   return true;
 }
 
-bool av1_compute_corner_list(const ImagePyramid *pyr, CornerList *corners) {
+bool av1_compute_corner_list(const YV12_BUFFER_CONFIG *frame, int bit_depth,
+                             int downsample_level, CornerList *corners) {
   assert(corners);
 
 #if CONFIG_MULTITHREAD
@@ -107,7 +126,8 @@ bool av1_compute_corner_list(const ImagePyramid *pyr, CornerList *corners) {
 #endif  // CONFIG_MULTITHREAD
 
   if (!corners->valid) {
-    corners->valid = compute_corner_list(pyr, corners);
+    corners->valid =
+        compute_corner_list(frame, bit_depth, downsample_level, corners);
   }
   bool valid = corners->valid;