diff --git a/src/encoder_state-ctors_dtors.c b/src/encoder_state-ctors_dtors.c index e53ac4e4..49d438b9 100644 --- a/src/encoder_state-ctors_dtors.c +++ b/src/encoder_state-ctors_dtors.c @@ -53,11 +53,16 @@ static int encoder_state_config_frame_init(encoder_state_t * const state) { state->frame->rc_alpha = 3.2003; state->frame->rc_beta = -1.367; + const encoder_control_t * const encoder = state->encoder_control; + const int num_lcus = encoder->in.width_in_lcu * encoder->in.height_in_lcu; + state->frame->lcu_stats = MALLOC(lcu_stats_t, num_lcus); + return 1; } static void encoder_state_config_frame_finalize(encoder_state_t * const state) { kvz_image_list_destroy(state->frame->ref); + FREE_POINTER(state->frame->lcu_stats); } static int encoder_state_config_tile_init(encoder_state_t * const state, diff --git a/src/encoderstate.c b/src/encoderstate.c index 31212ebe..967cfa8c 100644 --- a/src/encoderstate.c +++ b/src/encoderstate.c @@ -202,7 +202,7 @@ static void encoder_state_worker_encode_lcu(void * opaque) { const encoder_control_t * const encoder = state->encoder_control; videoframe_t* const frame = state->tile->frame; - kvz_set_lcu_lambda_and_qp(state); + kvz_set_lcu_lambda_and_qp(state, lcu->position); //This part doesn't write to bitstream, it's only search, deblock and sao @@ -241,6 +241,7 @@ static void encoder_state_worker_encode_lcu(void * opaque) { } //Now write data to bitstream (required to have a correct CABAC state) + const uint64_t existing_bits = kvz_bitstream_tell(&state->stream); //First LCU, and we are in a slice. We need a slice header if (state->type == ENCODER_STATE_TYPE_SLICE && lcu->index == 0) { @@ -266,7 +267,10 @@ static void encoder_state_worker_encode_lcu(void * opaque) { //Always 0 since otherwise it would be split kvz_cabac_encode_bin_trm(&state->cabac, 0); // end_of_slice_segment_flag } - + + const uint32_t bits = kvz_bitstream_tell(&state->stream) - existing_bits; + kvz_get_lcu_stats(state, lcu->position.x, lcu->position.y)->bits = bits; + //Wavefronts need the context to be copied to the next row if (state->type == ENCODER_STATE_TYPE_WAVEFRONT_ROW && lcu->index == 1) { int j; @@ -860,6 +864,22 @@ static void encoder_state_init_children(encoder_state_t * const state) { } } +static void normalize_lcu_weights(encoder_state_t * const state) +{ + if (state->frame->num == 0) return; + + const uint32_t num_lcus = state->encoder_control->in.width_in_lcu * + state->encoder_control->in.height_in_lcu; + double sum = 0.0; + for (uint32_t i = 0; i < num_lcus; i++) { + sum += state->frame->lcu_stats[i].weight; + } + + for (uint32_t i = 0; i < num_lcus; i++) { + state->frame->lcu_stats[i].weight /= sum; + } +} + static void encoder_state_init_new_frame(encoder_state_t * const state, kvz_picture* frame) { assert(state->type == ENCODER_STATE_TYPE_MAIN); @@ -908,6 +928,7 @@ static void encoder_state_init_new_frame(encoder_state_t * const state, kvz_pict encoder_state_remove_refs(state); encoder_state_ref_sort(state); + normalize_lcu_weights(state); kvz_set_picture_lambda_and_qp(state); encoder_state_init_children(state); @@ -1045,3 +1066,11 @@ coeff_scan_order_t kvz_get_scan_order(int8_t cu_type, int intra_mode, int depth) return SCAN_DIAG; } + +lcu_stats_t* kvz_get_lcu_stats(encoder_state_t *state, int lcu_x, int lcu_y) +{ + const int index = lcu_x + state->tile->lcu_offset_x + + (lcu_y + state->tile->lcu_offset_y) * + state->encoder_control->in.width_in_lcu; + return &state->frame->lcu_stats[index]; +} diff --git a/src/encoderstate.h b/src/encoderstate.h index ab3603b9..d883c11a 100644 --- a/src/encoderstate.h +++ b/src/encoderstate.h @@ -49,6 +49,23 @@ typedef enum { } encoder_state_type; +typedef struct lcu_stats_t { + //! \brief Number of bits that were spent + uint32_t bits; + + //! \brief Weight of the LCU for rate control + double weight; + + //! \brief Lambda value which was used for this LCU + double lambda; + + //! \brief Rate control alpha parameter + double rc_alpha; + + //! \brief Rate control beta parameter + double rc_beta; +} lcu_stats_t; + typedef struct encoder_state_config_frame_t { /** @@ -114,6 +131,13 @@ typedef struct encoder_state_config_frame_t { */ bool done; + /** + * \brief Information about the coded LCUs. + * + * Used for rate control. + */ + lcu_stats_t *lcu_stats; + } encoder_state_config_frame_t; typedef struct encoder_state_config_tile_t { @@ -249,6 +273,8 @@ void kvz_encoder_get_ref_lists(const encoder_state_t *const state, int ref_list_len_out[2], int ref_list_poc_out[2][16]); +lcu_stats_t* kvz_get_lcu_stats(encoder_state_t *state, int lcu_x, int lcu_y); + static const uint8_t g_group_idx[32] = { 0, 1, 2, 3, 4, 4, 5, 5, 6, 6, 6, 6, 7, 7, 7, 7, 8, 8, 8, 8, diff --git a/src/rate_control.c b/src/rate_control.c index 5be244ae..9ce91959 100644 --- a/src/rate_control.c +++ b/src/rate_control.c @@ -190,9 +190,15 @@ void kvz_set_picture_lambda_and_qp(encoder_state_t * const state) } } -void kvz_set_lcu_lambda_and_qp(encoder_state_t * const state) +void kvz_set_lcu_lambda_and_qp(encoder_state_t * const state, + vector2d_t pos) { state->lambda = state->frame->lambda; state->lambda_sqrt = sqrt(state->frame->lambda); state->qp = state->frame->QP; + + lcu_stats_t *lcu_stats = kvz_get_lcu_stats(state, pos.x, pos.y); + lcu_stats->lambda = state->lambda; + lcu_stats->rc_alpha = state->frame->rc_alpha; + lcu_stats->rc_beta = state->frame->rc_beta; } diff --git a/src/rate_control.h b/src/rate_control.h index b1e9281d..1ead1ca6 100644 --- a/src/rate_control.h +++ b/src/rate_control.h @@ -32,6 +32,7 @@ void kvz_set_picture_lambda_and_qp(encoder_state_t * const state); -void kvz_set_lcu_lambda_and_qp(encoder_state_t * const state); +void kvz_set_lcu_lambda_and_qp(encoder_state_t * const state, + vector2d_t pos); #endif // RATE_CONTROL_H_ diff --git a/src/search.c b/src/search.c index 440e7dca..6d31a50f 100644 --- a/src/search.c +++ b/src/search.c @@ -949,7 +949,10 @@ void kvz_search_lcu(encoder_state_t * const state, const int x, const int y, con } // Start search from depth 0. - search_cu(state, x, y, 0, work_tree); + double cost = search_cu(state, x, y, 0, work_tree); + + // Save squared cost for rate control. + kvz_get_lcu_stats(state, x / LCU_WIDTH, y / LCU_WIDTH)->weight = cost * cost; // The best decisions through out the LCU got propagated back to depth 0, // so copy those back to the frame.