340 lines
13 KiB
C++
340 lines
13 KiB
C++
|
/*
|
||
|
|
* Copyright (c) 2022, Alliance for Open Media. All rights reserved
|
||
|
|
*
|
||
|
|
* This source code is subject to the terms of the BSD 2 Clause License and
|
||
|
|
* the Alliance for Open Media Patent License 1.0. If the BSD 2 Clause License
|
||
|
|
* was not distributed with this source code in the LICENSE file, you can
|
||
|
|
* obtain it at www.aomedia.org/license/software. If the Alliance for Open
|
||
|
|
* Media Patent License 1.0 was not distributed with this source code in the
|
||
|
|
* PATENTS file, you can obtain it at www.aomedia.org/license/patent.
|
||
|
|
*/
|
||
|
|
|
||
|
|
#include <algorithm>
|
||
|
|
#include <set>
|
||
|
|
#include <utility>
|
||
|
|
#include <tuple>
|
||
|
|
#include <vector>
|
||
|
|
|
||
|
|
#include "av1/qmode_rc/reference_manager.h"
|
||
|
|
#include "av1/qmode_rc/ratectrl_qmode.h"
|
||
|
|
|
||
|
|
namespace aom {
|
||
|
|
|
||
|
|
void RefFrameManager::Reset() {
|
||
|
|
free_ref_idx_list_.clear();
|
||
|
|
for (int i = 0; i < static_cast<int>(ref_frame_table_.size()); ++i) {
|
||
|
|
free_ref_idx_list_.push_back(i);
|
||
|
|
ref_frame_table_[i] = GopFrameInvalid();
|
||
|
|
}
|
||
|
|
forward_stack_.clear();
|
||
|
|
backward_queue_.clear();
|
||
|
|
last_queue_.clear();
|
||
|
|
}
|
||
|
|
|
||
|
|
int RefFrameManager::AllocateRefIdx() {
|
||
|
|
if (free_ref_idx_list_.empty()) {
|
||
|
|
size_t backward_size = backward_queue_.size();
|
||
|
|
size_t last_size = last_queue_.size();
|
||
|
|
if (last_size >= backward_size) {
|
||
|
|
int ref_idx = last_queue_.front();
|
||
|
|
last_queue_.pop_front();
|
||
|
|
free_ref_idx_list_.push_back(ref_idx);
|
||
|
|
} else {
|
||
|
|
int ref_idx = backward_queue_.front();
|
||
|
|
backward_queue_.pop_front();
|
||
|
|
free_ref_idx_list_.push_back(ref_idx);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
int ref_idx = free_ref_idx_list_.front();
|
||
|
|
free_ref_idx_list_.pop_front();
|
||
|
|
return ref_idx;
|
||
|
|
}
|
||
|
|
|
||
|
|
int RefFrameManager::GetRefFrameCountByType(
|
||
|
|
RefUpdateType ref_update_type) const {
|
||
|
|
size_t cnt = 0;
|
||
|
|
switch (ref_update_type) {
|
||
|
|
case RefUpdateType::kForward: cnt = forward_stack_.size(); break;
|
||
|
|
case RefUpdateType::kBackward: cnt = backward_queue_.size(); break;
|
||
|
|
case RefUpdateType::kLast: cnt = last_queue_.size(); break;
|
||
|
|
case RefUpdateType::kNone: cnt = 0; break;
|
||
|
|
}
|
||
|
|
return static_cast<int>(cnt);
|
||
|
|
}
|
||
|
|
|
||
|
|
int RefFrameManager::GetRefFrameCount() const {
|
||
|
|
return GetRefFrameCountByType(RefUpdateType::kForward) +
|
||
|
|
GetRefFrameCountByType(RefUpdateType::kBackward) +
|
||
|
|
GetRefFrameCountByType(RefUpdateType::kLast);
|
||
|
|
}
|
||
|
|
|
||
|
|
// TODO(angiebird): Add unit test.
|
||
|
|
// Find the ref_idx corresponding to a ref_update_type.
|
||
|
|
// Return -1 if no ref frame is found.
|
||
|
|
// The priority_idx indicate closeness between the current frame and
|
||
|
|
// the ref frame in display order.
|
||
|
|
// For example, ref_update_type == kForward and priority_idx == 0 means
|
||
|
|
// find the closest ref frame in forward_stack_.
|
||
|
|
int RefFrameManager::GetRefFrameIdxByPriority(RefUpdateType ref_update_type,
|
||
|
|
int priority_idx) const {
|
||
|
|
if (ref_update_type == RefUpdateType::kForward) {
|
||
|
|
int size = static_cast<int>(forward_stack_.size());
|
||
|
|
// When two or more forward reference frames can be used, first get
|
||
|
|
// the highest quality one as the ARF, then going from nearest to
|
||
|
|
// the more distant ones in the forward reference frame list.
|
||
|
|
if (priority_idx < size) {
|
||
|
|
if (allow_two_fwd_frames_) {
|
||
|
|
if (priority_idx == 0) return forward_stack_[0];
|
||
|
|
return forward_stack_[size - priority_idx];
|
||
|
|
}
|
||
|
|
|
||
|
|
// Handle the special case where only one forward reference frame
|
||
|
|
// can be used. In this setting, we prefer the nearest frame.
|
||
|
|
return forward_stack_[size - 1 - priority_idx];
|
||
|
|
}
|
||
|
|
} else if (ref_update_type == RefUpdateType::kBackward) {
|
||
|
|
int size = static_cast<int>(backward_queue_.size());
|
||
|
|
if (priority_idx < size) {
|
||
|
|
return backward_queue_[size - priority_idx - 1];
|
||
|
|
}
|
||
|
|
} else if (ref_update_type == RefUpdateType::kLast) {
|
||
|
|
int size = static_cast<int>(last_queue_.size());
|
||
|
|
if (priority_idx < size) {
|
||
|
|
return last_queue_[size - priority_idx - 1];
|
||
|
|
}
|
||
|
|
}
|
||
|
|
return -1;
|
||
|
|
}
|
||
|
|
|
||
|
|
// The priority_idx indicate closeness between the current frame and
|
||
|
|
// the ref frame in display order.
|
||
|
|
// For example, ref_update_type == kForward and priority_idx == 0 means
|
||
|
|
// find the closest ref frame in forward_stack_.
|
||
|
|
GopFrame RefFrameManager::GetRefFrameByPriority(RefUpdateType ref_update_type,
|
||
|
|
int priority_idx) const {
|
||
|
|
int ref_idx = GetRefFrameIdxByPriority(ref_update_type, priority_idx);
|
||
|
|
if (ref_idx == -1) {
|
||
|
|
return GopFrameInvalid();
|
||
|
|
}
|
||
|
|
assert(ref_frame_table_[ref_idx].update_ref_idx == ref_idx);
|
||
|
|
return ref_frame_table_[ref_idx];
|
||
|
|
}
|
||
|
|
|
||
|
|
GopFrame RefFrameManager::GetRefFrameByIndex(int ref_idx) const {
|
||
|
|
return ref_frame_table_[ref_idx];
|
||
|
|
}
|
||
|
|
|
||
|
|
ReferenceName get_ref_name(RefUpdateType ref_update_type, int priority_idx,
|
||
|
|
const std::set<ReferenceName> &used_name_set) {
|
||
|
|
// TODO(angiebird): Find the better way to assign name lists.
|
||
|
|
// Maybe sort the names based on how frequent each name is being used in the
|
||
|
|
// past?
|
||
|
|
const std::vector<ReferenceName> forward_name_list{
|
||
|
|
ReferenceName::kAltrefFrame, ReferenceName::kBwdrefFrame,
|
||
|
|
ReferenceName::kAltref2Frame, ReferenceName::kGoldenFrame,
|
||
|
|
ReferenceName::kLast3Frame, ReferenceName::kLast2Frame,
|
||
|
|
ReferenceName::kLastFrame
|
||
|
|
};
|
||
|
|
const std::vector<ReferenceName> backward_name_list{
|
||
|
|
ReferenceName::kGoldenFrame, ReferenceName::kLastFrame,
|
||
|
|
ReferenceName::kLast2Frame, ReferenceName::kLast3Frame,
|
||
|
|
ReferenceName::kBwdrefFrame, ReferenceName::kAltref2Frame,
|
||
|
|
ReferenceName::kAltrefFrame
|
||
|
|
};
|
||
|
|
const std::vector<ReferenceName> last_name_list{
|
||
|
|
ReferenceName::kLastFrame, ReferenceName::kLast2Frame,
|
||
|
|
ReferenceName::kLast3Frame, ReferenceName::kGoldenFrame,
|
||
|
|
ReferenceName::kBwdrefFrame, ReferenceName::kAltref2Frame,
|
||
|
|
ReferenceName::kAltrefFrame
|
||
|
|
};
|
||
|
|
|
||
|
|
const std::vector<ReferenceName> *name_list = nullptr;
|
||
|
|
switch (ref_update_type) {
|
||
|
|
case RefUpdateType::kForward: name_list = &forward_name_list; break;
|
||
|
|
case RefUpdateType::kBackward: name_list = &backward_name_list; break;
|
||
|
|
case RefUpdateType::kLast: name_list = &last_name_list; break;
|
||
|
|
case RefUpdateType::kNone: break;
|
||
|
|
}
|
||
|
|
|
||
|
|
if (name_list) {
|
||
|
|
const int name_list_size = static_cast<int>(name_list->size());
|
||
|
|
for (int idx = priority_idx; idx < name_list_size; ++idx) {
|
||
|
|
ReferenceName ref_name = name_list->at(idx);
|
||
|
|
bool not_used = used_name_set.find(ref_name) == used_name_set.end();
|
||
|
|
if (not_used) return ref_name;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
return ReferenceName::kNoneFrame;
|
||
|
|
}
|
||
|
|
|
||
|
|
// Generate a list of available reference frames in priority order for the
|
||
|
|
// current to-be-coded frame. The list size should be less or equal to the size
|
||
|
|
// of ref_frame_table_. The reference frames with smaller indices are more
|
||
|
|
// likely to be a good reference frame. Therefore, they should be prioritized
|
||
|
|
// when the reference frame count is limited. For example, if we plan to use 3
|
||
|
|
// reference frames, we should choose ref_frame_list[0], ref_frame_list[1] and
|
||
|
|
// ref_frame_list[2].
|
||
|
|
std::vector<ReferenceFrame> RefFrameManager::GetRefFrameListByPriority() const {
|
||
|
|
constexpr int round_robin_size = 3;
|
||
|
|
const std::vector<RefUpdateType> round_robin_list{ RefUpdateType::kForward,
|
||
|
|
RefUpdateType::kBackward,
|
||
|
|
RefUpdateType::kLast };
|
||
|
|
std::vector<int> priority_idx_list(round_robin_size, 0);
|
||
|
|
int available_ref_frames = GetRefFrameCount();
|
||
|
|
std::vector<ReferenceFrame> ref_frame_list;
|
||
|
|
int ref_frame_count = 0;
|
||
|
|
int round_robin_idx = 0;
|
||
|
|
|
||
|
|
std::set<ReferenceName> used_name_set;
|
||
|
|
while (ref_frame_count < available_ref_frames &&
|
||
|
|
ref_frame_count < max_ref_frames_) {
|
||
|
|
const RefUpdateType ref_update_type = round_robin_list[round_robin_idx];
|
||
|
|
int priority_idx = priority_idx_list[round_robin_idx];
|
||
|
|
int ref_idx = GetRefFrameIdxByPriority(ref_update_type, priority_idx);
|
||
|
|
if (ref_idx != -1) {
|
||
|
|
const ReferenceName name =
|
||
|
|
get_ref_name(ref_update_type, priority_idx, used_name_set);
|
||
|
|
assert(name != ReferenceName::kNoneFrame);
|
||
|
|
used_name_set.insert(name);
|
||
|
|
ReferenceFrame ref_frame = { ref_idx, name };
|
||
|
|
ref_frame_list.push_back(ref_frame);
|
||
|
|
++ref_frame_count;
|
||
|
|
++priority_idx_list[round_robin_idx];
|
||
|
|
}
|
||
|
|
round_robin_idx = (round_robin_idx + 1) % round_robin_size;
|
||
|
|
}
|
||
|
|
return ref_frame_list;
|
||
|
|
}
|
||
|
|
|
||
|
|
void RefFrameManager::UpdateOrder(int global_order_idx) {
|
||
|
|
cur_global_order_idx_ = global_order_idx;
|
||
|
|
if (forward_stack_.empty()) {
|
||
|
|
return;
|
||
|
|
}
|
||
|
|
int ref_idx = forward_stack_.back();
|
||
|
|
const GopFrame &gf_frame = ref_frame_table_[ref_idx];
|
||
|
|
|
||
|
|
// If the current processing frame is an overlay / show existing frame.
|
||
|
|
if (gf_frame.global_order_idx == global_order_idx) {
|
||
|
|
forward_stack_.pop_back();
|
||
|
|
if (gf_frame.is_golden_frame) {
|
||
|
|
// high quality frame
|
||
|
|
backward_queue_.push_back(ref_idx);
|
||
|
|
} else {
|
||
|
|
last_queue_.push_back(ref_idx);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
int RefFrameManager::ColocatedRefIdx(int global_order_idx) {
|
||
|
|
if (forward_stack_.empty()) return -1;
|
||
|
|
int ref_idx = forward_stack_.back();
|
||
|
|
int arf_global_order_idx = ref_frame_table_[ref_idx].global_order_idx;
|
||
|
|
if (arf_global_order_idx == global_order_idx) {
|
||
|
|
return ref_idx;
|
||
|
|
}
|
||
|
|
return -1;
|
||
|
|
}
|
||
|
|
|
||
|
|
static RefUpdateType infer_ref_update_type(const GopFrame &gop_frame,
|
||
|
|
int cur_global_order_idx) {
|
||
|
|
if (gop_frame.global_order_idx > cur_global_order_idx) {
|
||
|
|
return RefUpdateType::kForward;
|
||
|
|
}
|
||
|
|
if (gop_frame.is_golden_frame) {
|
||
|
|
return RefUpdateType::kBackward;
|
||
|
|
}
|
||
|
|
if (gop_frame.encode_ref_mode == EncodeRefMode::kShowExisting ||
|
||
|
|
gop_frame.encode_ref_mode == EncodeRefMode::kOverlay) {
|
||
|
|
return RefUpdateType::kNone;
|
||
|
|
}
|
||
|
|
return RefUpdateType::kLast;
|
||
|
|
}
|
||
|
|
|
||
|
|
using PrimaryRefKey = std::tuple<int, // abs layer_depth delta
|
||
|
|
bool, // is_key_frame differs
|
||
|
|
bool, // is_golden_frame differs
|
||
|
|
bool, // is_arf_frame differs
|
||
|
|
bool, // is_show_frame differs
|
||
|
|
bool, // encode_ref_mode differs
|
||
|
|
int>; // abs order_idx delta
|
||
|
|
|
||
|
|
// Generate PrimaryRefKey based on abs layer_depth delta,
|
||
|
|
// frame flags and abs order_idx delta. These are the fields that will
|
||
|
|
// be used to pick the primary reference frame for probability model
|
||
|
|
static PrimaryRefKey get_primary_ref_key(const GopFrame &cur_frame,
|
||
|
|
const GopFrame &ref_frame) {
|
||
|
|
return std::make_tuple(abs(cur_frame.layer_depth - ref_frame.layer_depth),
|
||
|
|
cur_frame.is_key_frame != ref_frame.is_key_frame,
|
||
|
|
cur_frame.is_golden_frame != ref_frame.is_golden_frame,
|
||
|
|
cur_frame.is_arf_frame != ref_frame.is_arf_frame,
|
||
|
|
cur_frame.is_show_frame != ref_frame.is_show_frame,
|
||
|
|
cur_frame.encode_ref_mode != ref_frame.encode_ref_mode,
|
||
|
|
abs(cur_frame.order_idx - ref_frame.order_idx));
|
||
|
|
}
|
||
|
|
|
||
|
|
// Pick primary_ref_idx for probability model.
|
||
|
|
ReferenceFrame RefFrameManager::GetPrimaryRefFrame(
|
||
|
|
const GopFrame &gop_frame) const {
|
||
|
|
assert(gop_frame.is_valid);
|
||
|
|
std::vector<std::pair<PrimaryRefKey, int>> candidate_list;
|
||
|
|
for (auto &ref_frame_in_gop_frame : gop_frame.ref_frame_list) {
|
||
|
|
const GopFrame &ref_frame = ref_frame_table_[ref_frame_in_gop_frame.index];
|
||
|
|
if (ref_frame.is_valid) {
|
||
|
|
assert(ref_frame_in_gop_frame.index == ref_frame.update_ref_idx);
|
||
|
|
PrimaryRefKey key = get_primary_ref_key(gop_frame, ref_frame);
|
||
|
|
std::pair<PrimaryRefKey, int> candidate = {
|
||
|
|
key, ref_frame_in_gop_frame.index
|
||
|
|
};
|
||
|
|
candidate_list.push_back(candidate);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
std::sort(candidate_list.begin(), candidate_list.end());
|
||
|
|
|
||
|
|
ReferenceFrame ref_frame = { -1, ReferenceName::kNoneFrame };
|
||
|
|
assert(candidate_list.size() == gop_frame.ref_frame_list.size());
|
||
|
|
if (!candidate_list.empty()) {
|
||
|
|
int ref_idx = candidate_list[0].second;
|
||
|
|
for (const auto &frame : gop_frame.ref_frame_list) {
|
||
|
|
if (frame.index == ref_idx) {
|
||
|
|
ref_frame = frame;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
return ref_frame;
|
||
|
|
}
|
||
|
|
|
||
|
|
void RefFrameManager::UpdateRefFrameTable(GopFrame *gop_frame) {
|
||
|
|
allow_two_fwd_frames_ =
|
||
|
|
(max_ref_frames_ - !!GetRefFrameCountByType(RefUpdateType::kBackward) -
|
||
|
|
!!GetRefFrameCountByType(RefUpdateType::kLast)) >= 2;
|
||
|
|
gop_frame->ref_frame_list = GetRefFrameListByPriority();
|
||
|
|
gop_frame->primary_ref_frame = GetPrimaryRefFrame(*gop_frame);
|
||
|
|
gop_frame->colocated_ref_idx = ColocatedRefIdx(gop_frame->global_order_idx);
|
||
|
|
|
||
|
|
if (gop_frame->is_show_frame) {
|
||
|
|
UpdateOrder(gop_frame->global_order_idx);
|
||
|
|
}
|
||
|
|
// Call infer_ref_update_type() after UpdateOrder() so that
|
||
|
|
// cur_global_order_idx_ is up-to-date
|
||
|
|
RefUpdateType ref_update_type =
|
||
|
|
infer_ref_update_type(*gop_frame, cur_global_order_idx_);
|
||
|
|
if (ref_update_type == RefUpdateType::kNone) {
|
||
|
|
gop_frame->update_ref_idx = -1;
|
||
|
|
} else {
|
||
|
|
const int ref_idx = AllocateRefIdx();
|
||
|
|
gop_frame->update_ref_idx = ref_idx;
|
||
|
|
switch (ref_update_type) {
|
||
|
|
case RefUpdateType::kForward: forward_stack_.push_back(ref_idx); break;
|
||
|
|
case RefUpdateType::kBackward: backward_queue_.push_back(ref_idx); break;
|
||
|
|
case RefUpdateType::kLast: last_queue_.push_back(ref_idx); break;
|
||
|
|
case RefUpdateType::kNone: break;
|
||
|
|
}
|
||
|
|
ref_frame_table_[ref_idx] = *gop_frame;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
} // namespace aom
|