Skip to content

Commit

Permalink
Add support for recon gain parameters in the ParametersManager.
Browse files Browse the repository at this point in the history 
This CL adds the requisite member fields needed to keep track of recon gain parameters separately from demixing parameters, as well as initializes these fields in Initialize().

PiperOrigin-RevId: 656489728
  • Loading branch information
@jwcullen
Googler authored and jwcullen committed Jul 29, 2024
1 parent 63c0a9f commit 75855b5
Show file tree
Hide file tree
Showing 4 changed files with 142 additions and 35 deletions.
74 changes: 48 additions & 26 deletions iamf/cli/parameters_manager.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -33,34 +33,43 @@ ParametersManager::ParametersManager(
: audio_elements_(audio_elements) {}

absl::Status ParametersManager::Initialize() {
// Collect all `DemixingParamDefinition`s in all Audio Elements. Validate
// there is no more than one per Audio Element.
// Collect all `DemixingParamDefinition`s and all `ReconGainParamDefinitions`
// in all Audio Elements. Validate there is no more than one per Audio
// Element.
for (const auto& [audio_element_id, audio_element] : audio_elements_) {
const DemixingParamDefinition* demixing_param_definition = nullptr;
const ReconGainParamDefinition* recon_gain_param_definition = nullptr;
for (const auto& param : audio_element.obu.audio_element_params_) {
if (param.param_definition_type !=
if (param.param_definition_type ==
ParamDefinition::kParameterDefinitionDemixing) {
continue;
if (demixing_param_definition != nullptr) {
return absl::InvalidArgumentError(
"Not allowed to have multiple demixing parameters in a "
"single Audio Element.");
}

demixing_param_definition =
static_cast<DemixingParamDefinition*>(param.param_definition.get());

// Continue searching. Only to validate that there is at most one
// `DemixingParamDefinition`.
} else if (param.param_definition_type ==
ParamDefinition::kParameterDefinitionReconGain) {
if (recon_gain_param_definition != nullptr) {
return absl::InvalidArgumentError(
"Not allowed to have multiple recon gain parameters in a "
"single Audio Element.");
}
recon_gain_param_definition = static_cast<ReconGainParamDefinition*>(
param.param_definition.get());
}

if (demixing_param_definition != nullptr) {
return absl::InvalidArgumentError(
"Not allowed to have multiple demixing parameters in a "
"single Audio Element.");
}

demixing_param_definition =
static_cast<DemixingParamDefinition*>(param.param_definition.get());

// Continue searching. Only to validate that there is at most one
// `DemixingParamDefinition`.
}

if (demixing_param_definition != nullptr) {
// Insert a `nullptr` for a parameter ID. If no parameter blocks have
// this parameter ID, then it will remain null and default values will
// be used.
parameter_blocks_.insert(
demixing_parameter_blocks_.insert(
{demixing_param_definition->parameter_id_, nullptr});
demixing_states_[audio_element_id] = {
.param_definition = demixing_param_definition,
Expand All @@ -69,6 +78,17 @@ absl::Status ParametersManager::Initialize() {
.update_rule = DemixingInfoParameterData::kFirstFrame,
};
}
if (recon_gain_param_definition != nullptr) {
// Insert a `nullptr` for a parameter ID. If no parameter blocks have
// this parameter ID, then it will remain null and default values will
// be used.
recon_gain_parameter_blocks_.insert(
{recon_gain_param_definition->parameter_id_, nullptr});
recon_gain_states_[audio_element_id] = {
.param_definition = recon_gain_param_definition,
.next_timestamp = 0,
};
}
}

return absl::OkStatus();
Expand Down Expand Up @@ -96,11 +116,11 @@ absl::Status ParametersManager::GetDownMixingParameters(
auto& demixing_state = demixing_states_iter->second;
const auto* param_definition = demixing_state.param_definition;
const auto* parameter_block =
parameter_blocks_.at(param_definition->parameter_id_);
demixing_parameter_blocks_.at(param_definition->parameter_id_);
if (parameter_block == nullptr) {
// Failed to find a parameter block that overlaps this frame. Use the
// default value from the parameter definition. This is OK when there are no
// parameter blocks covering this substream. If there is only partial
// default value from the parameter definition. This is OK when there are
// no parameter blocks covering this substream. If there is only partial
// coverage this will be marked invalid when the coverage of parameter
// blocks is checked.
LOG_FIRST_N(WARNING, 10)
Expand All @@ -114,8 +134,9 @@ absl::Status ParametersManager::GetDownMixingParameters(

if (parameter_block->start_timestamp != demixing_state.next_timestamp) {
return absl::InvalidArgumentError(absl::StrCat(
"Mismatching timestamps for down-mixing parameters for audio element ",
"ID= ", audio_element_id, ": expecting", demixing_state.next_timestamp,
"Mismatching timestamps for down-mixing parameters for "
"audio element ID= ",
audio_element_id, ": expecting", demixing_state.next_timestamp,
" but got ", parameter_block->start_timestamp));
}

Expand All @@ -131,7 +152,8 @@ absl::Status ParametersManager::GetDownMixingParameters(

void ParametersManager::AddDemixingParameterBlock(
const ParameterBlockWithData* parameter_block) {
parameter_blocks_[parameter_block->obu->parameter_id_] = parameter_block;
demixing_parameter_blocks_[parameter_block->obu->parameter_id_] =
parameter_block;
}

absl::Status ParametersManager::UpdateDemixingState(
Expand All @@ -147,11 +169,11 @@ absl::Status ParametersManager::UpdateDemixingState(
auto& demixing_state = demixing_states_iter->second;

// Using `.at()` here is safe because if the demixing state exists for the
// `audio_element_id`, an entry in `parameter_blocks_` with the key
// `audio_element_id`, an entry in `demixing_parameter_blocks_` with the key
// `demixing_state.param_definition->parameter_id_` has already been
// created during `Initialize()`.
auto& parameter_block =
parameter_blocks_.at(demixing_state.param_definition->parameter_id_);
auto& parameter_block = demixing_parameter_blocks_.at(
demixing_state.param_definition->parameter_id_);
if (parameter_block == nullptr) {
// No parameter block found for this ID. Do not validate the timestamp
// or update anything else.
Expand Down
18 changes: 16 additions & 2 deletions iamf/cli/parameters_manager.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -109,16 +109,30 @@ class ParametersManager {
DemixingInfoParameterData::WIdxUpdateRule update_rule;
};

struct ReconGainState {
const ReconGainParamDefinition* param_definition;

// Timestamp for the next frame to be processed.
int32_t next_timestamp;
};

// Mapping from Audio Element ID to audio element data.
const absl::flat_hash_map<DecodedUleb128, AudioElementWithData>&
audio_elements_;

// Mapping from Parameter ID to parameter blocks.
// Mapping from Parameter ID to demixing parameter blocks.
absl::flat_hash_map<DecodedUleb128, const ParameterBlockWithData*>
demixing_parameter_blocks_;

// Mapping from Parameter ID to recon gain parameter blocks.
absl::flat_hash_map<DecodedUleb128, const ParameterBlockWithData*>
parameter_blocks_;
recon_gain_parameter_blocks_;

// Mapping from Audio Element ID to the demixing state.
absl::flat_hash_map<DecodedUleb128, DemixingState> demixing_states_;

// Mapping from Audio Element ID to the recon gain state.
absl::flat_hash_map<DecodedUleb128, ReconGainState> recon_gain_states_;
};

} // namespace iamf_tools
Expand Down
1 change: 1 addition & 0 deletions iamf/cli/tests/BUILD
View file
Original file line number Diff line number Diff line change
Expand Up @@ -285,6 +285,7 @@ cc_test(
"//iamf/cli:audio_element_with_data",
"//iamf/cli:parameter_block_with_data",
"//iamf/cli:parameters_manager",
"//iamf/obu:audio_element",
"//iamf/obu:codec_config",
"//iamf/obu:demixing_info_param_data",
"//iamf/obu:leb128",
Expand Down
84 changes: 77 additions & 7 deletions iamf/cli/tests/parameters_manager_test.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@

#include <cstdint>
#include <memory>
#include <utility>
#include <vector>

#include "absl/container/flat_hash_map.h"
Expand All @@ -23,6 +24,7 @@
#include "iamf/cli/audio_element_with_data.h"
#include "iamf/cli/parameter_block_with_data.h"
#include "iamf/cli/tests/cli_test_utils.h"
#include "iamf/obu/audio_element.h"
#include "iamf/obu/codec_config.h"
#include "iamf/obu/demixing_info_param_data.h"
#include "iamf/obu/leb128.h"
Expand All @@ -39,18 +41,14 @@ constexpr DecodedUleb128 kCodecConfigId = 1450;
constexpr DecodedUleb128 kSampleRate = 16000;
constexpr DecodedUleb128 kAudioElementId = 157;
constexpr DecodedUleb128 kParameterId = 995;
constexpr DecodedUleb128 kSecondParameterId = 996;
constexpr DecodedUleb128 kDuration = 8;
constexpr DemixingInfoParameterData::DMixPMode kDMixPMode =
DemixingInfoParameterData::kDMixPMode3_n;

absl::Status AddOneDemixingParameterBlock(
const ParamDefinition& param_definition, int32_t start_timestamp,
PerIdParameterMetadata& per_id_metadata,
absl::Status AppendParameterBlock(
int32_t start_timestamp, PerIdParameterMetadata& per_id_metadata,
std::vector<ParameterBlockWithData>& parameter_blocks) {
per_id_metadata = {
.param_definition_type = ParamDefinition::kParameterDefinitionDemixing,
.param_definition = param_definition,
};
parameter_blocks.emplace_back(ParameterBlockWithData{
std::make_unique<ParameterBlockObu>(ObuHeader(), kParameterId,
per_id_metadata),
Expand All @@ -59,13 +57,72 @@ absl::Status AddOneDemixingParameterBlock(
absl::Status status =
parameter_block_obu.InitializeSubblocks(kDuration, kDuration, 1);
status.Update(parameter_block_obu.SetSubblockDuration(0, kDuration));
return status;
}

absl::Status AddOneDemixingParameterBlock(
const ParamDefinition& param_definition, int32_t start_timestamp,
PerIdParameterMetadata& per_id_metadata,
std::vector<ParameterBlockWithData>& parameter_blocks) {
per_id_metadata = {
.param_definition_type = ParamDefinition::kParameterDefinitionDemixing,
.param_definition = param_definition,
};
auto status =
AppendParameterBlock(start_timestamp, per_id_metadata, parameter_blocks);
DemixingInfoParameterData demixing_info_param_data;
demixing_info_param_data.dmixp_mode = kDMixPMode;
ParameterBlockObu& parameter_block_obu = *parameter_blocks.back().obu;
parameter_block_obu.subblocks_[0].param_data = demixing_info_param_data;

return status;
}

absl::Status AddOneReconGainParameterBlock(
const ParamDefinition& param_definition, int32_t start_timestamp,
PerIdParameterMetadata& per_id_metadata,
std::vector<ParameterBlockWithData>& parameter_blocks) {
per_id_metadata = {
.param_definition_type = ParamDefinition::kParameterDefinitionReconGain,
.param_definition = param_definition,
};
auto status =
AppendParameterBlock(start_timestamp, per_id_metadata, parameter_blocks);

ReconGainInfoParameterData recon_gain_info_param_data;
recon_gain_info_param_data.recon_gain_elements.emplace_back(ReconGainElement{
.recon_gain_flag = DecodedUleb128(1),
.recon_gain = {0},
});
ParameterBlockObu& parameter_block_obu = *parameter_blocks.back().obu;
parameter_block_obu.subblocks_[0].param_data = recon_gain_info_param_data;

return status;
}

// TODO(b/355436892): Refactor common parts of this and
// `AddDemixingParamDefinition()` into a helper function.
void AddReconGainParamDefinition(DecodedUleb128 parameter_id,
DecodedUleb128 parameter_rate,
DecodedUleb128 duration,
AudioElementObu& audio_element_obu) {
const auto audio_element_id = audio_element_obu.GetAudioElementId();
auto param_definition =
std::make_unique<ReconGainParamDefinition>(audio_element_id);
param_definition->parameter_id_ = parameter_id;
param_definition->parameter_rate_ = parameter_rate;
param_definition->param_definition_mode_ = 0;
param_definition->reserved_ = 0;
param_definition->duration_ = duration;
param_definition->constant_subblock_duration_ = duration;

// Add to the Audio Element OBU.
audio_element_obu.InitializeParams(audio_element_obu.num_parameters_ + 1);
audio_element_obu.audio_element_params_.back() = AudioElementParam{
.param_definition_type = ParamDefinition::kParameterDefinitionReconGain,
.param_definition = std::move(param_definition)};
}

class ParametersManagerTest : public testing::Test {
public:
ParametersManagerTest() {
Expand Down Expand Up @@ -110,6 +167,19 @@ TEST_F(ParametersManagerTest, InitializeWithTwoDemixingParametersFails) {
EXPECT_FALSE(parameters_manager_->Initialize().ok());
}

TEST_F(ParametersManagerTest, InitializeWithReconGainParameterSucceeds) {
AddReconGainParamDefinition(kSecondParameterId, kSampleRate, kDuration,
audio_elements_.at(kAudioElementId).obu);
EXPECT_THAT(AddOneReconGainParameterBlock(
*audio_elements_.at(kAudioElementId)
.obu.audio_element_params_[0]
.param_definition,
/*start_timestamp=*/0, per_id_metadata_, parameter_blocks_),
IsOk());
parameters_manager_ = std::make_unique<ParametersManager>(audio_elements_);
EXPECT_THAT(parameters_manager_->Initialize(), IsOk());
}

TEST_F(ParametersManagerTest, DemixingParamDefinitionIsAvailable) {
parameters_manager_ = std::make_unique<ParametersManager>(audio_elements_);
ASSERT_THAT(parameters_manager_->Initialize(), IsOk());
Expand Down

0 comments on commit 75855b5

Please sign in to comment.