Add psnr and sse to EncodeFrameResult

Change-Id: I33c410a14b86f95278eff8d1d0e6992f1b82a17d

diff --git a/test/simple_encode_test.cc b/test/simple_encode_test.cc
index 30970514f3865dbe18c26bff18769581ce8a15b8..3184d00588d1f8b9d151a7cea08d65e7b38291fa 100644 (file)
--- a/test/simple_encode_test.cc
+++ b/test/simple_encode_test.cc
@@ -54,7 +54,7 @@ TEST(SimpleEncode, EncodeFrame) {
   int h = 288;
   int frame_rate_num = 30;
   int frame_rate_den = 1;
-  int target_bitrate = 200;
+  int target_bitrate = 1000;
   int num_frames = 17;
   // TODO(angiebird): Figure out how to upload test video to our codebase
   FILE *file = fopen("bus_352x288_420_f20_b8.yuv", "r");
@@ -83,8 +83,11 @@ TEST(SimpleEncode, EncodeFrame) {
     EXPECT_LT(encode_frame_result.show_idx, num_frames);
     if (i == num_coding_frames - 1) {
       EXPECT_EQ(encode_frame_result.show_idx, num_frames - 1)
-          << "The last coding frame should should be the last display order";
+          << "The last coding frame should be the last display order";
     }
+    EXPECT_GE(encode_frame_result.psnr, 34)
+        << "The psnr is supposed to be greater than 34 given the "
+           "target_bitrate 1000 kbps";
   }
   EXPECT_EQ(num_alternate_refereces, ref_num_alternate_refereces);
   simple_encode.EndEncode();

diff --git a/vp9/encoder/vp9_encoder.c b/vp9/encoder/vp9_encoder.c
index 07f2d3444fde1f6b46354f9d5c5ad3b494596077..c080ad47dc1ca9109130f56f0499dc486f839085 100644 (file)
--- a/vp9/encoder/vp9_encoder.c
+++ b/vp9/encoder/vp9_encoder.c
@@ -5374,8 +5374,6 @@ static void Pass2Encode(VP9_COMP *cpi, size_t *size, uint8_t *dest,
   mismatch_move_frame_idx_w();
 #endif
   encode_frame_to_data_rate(cpi, size, dest, frame_flags);
-
-  vp9_twopass_postencode_update(cpi);
 }
 #endif  // !CONFIG_REALTIME_ONLY
 
@@ -7091,12 +7089,29 @@ static void init_encode_frame_result(ENCODE_FRAME_RESULT *encode_frame_result) {
   encode_frame_result->show_idx = -1;  // Actual encoding deosn't happen.
 }
 
+#if !CONFIG_REALTIME_ONLY
 static void update_encode_frame_result(ENCODE_FRAME_RESULT *encode_frame_result,
                                        int show_idx,
-                                       FRAME_UPDATE_TYPE update_type) {
+                                       FRAME_UPDATE_TYPE update_type,
+                                       const YV12_BUFFER_CONFIG *source_frame,
+                                       const YV12_BUFFER_CONFIG *coded_frame,
+                                       uint32_t bit_depth,
+                                       uint32_t input_bit_depth) {
+  PSNR_STATS psnr;
+#if CONFIG_VP9_HIGHBITDEPTH
+  vpx_calc_highbd_psnr(source_frame, coded_frame, &psnr, bit_depth,
+                       input_bit_depth);
+#else
+  (void)bit_depth;
+  (void)input_bit_depth;
+  vpx_calc_psnr(source_frame, coded_frame, &psnr);
+#endif
+  encode_frame_result->psnr = psnr.psnr[0];
+  encode_frame_result->sse = psnr.sse[0];
   encode_frame_result->show_idx = show_idx;
   encode_frame_result->update_type = update_type;
 }
+#endif  // !CONFIG_REALTIME_ONLY
 
 int vp9_get_compressed_data(VP9_COMP *cpi, unsigned int *frame_flags,
                             size_t *size, uint8_t *dest, int64_t *time_stamp,
@@ -7349,12 +7364,6 @@ int vp9_get_compressed_data(VP9_COMP *cpi, unsigned int *frame_flags,
   bitstream_queue_set_frame_write(cm->current_video_frame * 2 + cm->show_frame);
 #endif
 
-  if (oxcf->pass != 1) {
-    update_encode_frame_result(
-        encode_frame_result, source->show_idx,
-        cpi->twopass.gf_group.update_type[cpi->twopass.gf_group.index]);
-  }
-
   cpi->td.mb.fp_src_pred = 0;
 #if CONFIG_REALTIME_ONLY
   if (cpi->use_svc) {
@@ -7381,6 +7390,25 @@ int vp9_get_compressed_data(VP9_COMP *cpi, unsigned int *frame_flags,
     vp9_first_pass(cpi, source);
   } else if (oxcf->pass == 2 && !cpi->use_svc) {
     Pass2Encode(cpi, size, dest, frame_flags);
+    // update_encode_frame_result() depends on twopass.gf_group.index and
+    // cm->new_fb_idx and cpi->Source are updated for current properly and have
+    // not been updated for the next frame yet.
+    // The update locations are as follows.
+    // 1) twopass.gf_group.index is initialized at define_gf_group by vp9_zero()
+    // for the first frame in the gf_group and is updated for the next frame at
+    // vp9_twopass_postencode_update().
+    // 2) cpi->Source is updated at the beginging of this function, i.e.
+    // vp9_get_compressed_data()
+    // 3) cm->new_fb_idx is updated at the beginging of this function by
+    // get_free_fb(cm)
+    // TODO(angiebird): Improve the codebase to make the update of frame
+    // dependent variables more robust.
+    update_encode_frame_result(
+        encode_frame_result, source->show_idx,
+        cpi->twopass.gf_group.update_type[cpi->twopass.gf_group.index],
+        cpi->Source, get_frame_new_buffer(cm), cpi->oxcf.input_bit_depth,
+        cm->bit_depth);
+    vp9_twopass_postencode_update(cpi);
   } else if (cpi->use_svc) {
     SvcEncode(cpi, size, dest, frame_flags);
   } else {

diff --git a/vp9/encoder/vp9_encoder.h b/vp9/encoder/vp9_encoder.h
index 2304209531be85d48790f11f67e62920d4fd86fa..66526bd00df0412fd80e7fd81ad35a3a6e3ed645 100644 (file)
--- a/vp9/encoder/vp9_encoder.h
+++ b/vp9/encoder/vp9_encoder.h
@@ -825,6 +825,8 @@ typedef struct VP9_COMP {
 typedef struct ENCODE_FRAME_RESULT {
   int show_idx;
   FRAME_UPDATE_TYPE update_type;
+  double psnr;
+  uint64_t sse;
 } ENCODE_FRAME_RESULT;
 
 void vp9_initialize_enc(void);

diff --git a/vp9/simple_encode.cc b/vp9/simple_encode.cc
index 9319a7f9fb00920de3fd299a180907c7f7b2f9df..4d7ec3084c9207aa0156934c89114a2922c170a5 100644 (file)
--- a/vp9/simple_encode.cc
+++ b/vp9/simple_encode.cc
@@ -90,6 +90,8 @@ static void update_encode_frame_result(
   encode_frame_result->show_idx = encode_frame_info->show_idx;
   encode_frame_result->frame_type =
       get_frame_type_from_update_type(encode_frame_info->update_type);
+  encode_frame_result->psnr = encode_frame_info->psnr;
+  encode_frame_result->sse = encode_frame_info->sse;
 }
 
 SimpleEncode::SimpleEncode(int frame_width, int frame_height,

diff --git a/vp9/simple_encode.h b/vp9/simple_encode.h
index 26415a3c2d99edf4cfdca5c896fd65289c540796..fa1847161d1354cfab050a81aa380349ce1211dd 100644 (file)
--- a/vp9/simple_encode.h
+++ b/vp9/simple_encode.h
@@ -14,7 +14,8 @@ struct EncodeFrameResult {
   // The EncodeFrame will allocate a buffer, write the coding data into the
   // buffer and give the ownership of the buffer to coding_data
   std::unique_ptr<unsigned char[]> coding_data;
-  // TODO(angiebird): double psnr;
+  double psnr;
+  uint64_t sse;
   // TODO(angiebird): int quantize_index ;
 };