diff --git a/iamf/cli/obu_sequencer.cc b/iamf/cli/obu_sequencer.cc index 56ddd4b..96970df 100644 --- a/iamf/cli/obu_sequencer.cc +++ b/iamf/cli/obu_sequencer.cc @@ -13,9 +13,14 @@ #include #include +#include +#include #include +#include #include #include +#include +#include #include #include @@ -44,6 +49,9 @@ namespace iamf_tools { namespace { +// Start with 64 KB. +constexpr int64_t kBufferStartSize = 65536; + template std::vector SortedKeys(const KeyValueMap& map, const KeyComparator& comparator) { @@ -56,6 +64,48 @@ std::vector SortedKeys(const KeyValueMap& map, return keys; } +// The caller of this function should pre-seed the write buffer with Descriptor +// OBUs. +absl::Status WriteIaSequenceToFile(const std::string& iamf_filename, + bool include_temporal_delimiters, + const TemporalUnitMap& temporal_unit_map, + WriteBitBuffer& wb) { + std::optional output_iamf; + if (!iamf_filename.empty()) { + output_iamf.emplace(iamf_filename, std::fstream::out | std::ios::binary); + } + + // Write all Audio Frame and Parameter Block OBUs ordered by temporal unit. + int num_samples = 0; + for (const auto& temporal_unit : temporal_unit_map) { + // The temporal units will typically be the largest part of an IAMF + // sequence. Occasionally flush to buffer to avoid keeping it all in memory. + RETURN_IF_NOT_OK(wb.MaybeFlushIfCloseToCapacity(output_iamf)); + + RETURN_IF_NOT_OK(ObuSequencerBase::WriteTemporalUnit( + include_temporal_delimiters, temporal_unit.second, wb, num_samples)); + } + LOG(INFO) << "Wrote " << temporal_unit_map.size() + << " temporal units with a total of " << num_samples + << " samples excluding padding."; + + // Flush any unwritten bytes before exiting. + RETURN_IF_NOT_OK(wb.FlushAndWriteToFile(output_iamf)); + return absl::OkStatus(); +} + +void MaybeRemoveFile(const std::string& filename) { + if (filename.empty()) { + return; + } + std::error_code error_code; + std::filesystem::remove(filename, error_code); + if (!error_code) { + // File clean up failed somehow. Just log the error and move on. + LOG(ERROR).WithPerror() << "Failed to remove " << filename; + } +} + } // namespace absl::Status ObuSequencerBase::GenerateTemporalUnitMap( @@ -290,10 +340,9 @@ absl::Status ObuSequencerIamf::PickAndPlace( const std::list& audio_frames, const std::list& parameter_blocks, const std::list& arbitrary_obus) { - // Write buffer. Let's start with 64 KB. The buffer will resize for larger - // OBUs if needed. - static const int64_t kBufferSize = 65536; - WriteBitBuffer wb(kBufferSize, leb_generator_); + // Seed with a reasonable starting size. It is arbitrary because + // `WriteBitBuffer`s automatically resize as needed. + WriteBitBuffer wb(kBufferStartSize, leb_generator_); RETURN_IF_NOT_OK(ArbitraryObu::WriteObusWithHook( ArbitraryObu::kInsertionHookBeforeDescriptors, arbitrary_obus, wb)); @@ -311,24 +360,12 @@ absl::Status ObuSequencerIamf::PickAndPlace( RETURN_IF_NOT_OK(ObuSequencerBase::GenerateTemporalUnitMap( audio_frames, parameter_blocks, arbitrary_obus, temporal_unit_map)); - // Write all Audio Frame and Parameter Block OBUs ordered by temporal unit. - int num_samples = 0; - for (const auto& temporal_unit : temporal_unit_map) { - // The temporal units will typically be the largest part of an IAMF - // sequence. Occasionally flush to buffer to avoid keeping it all in memory. - RETURN_IF_NOT_OK(wb.MaybeFlushIfCloseToCapacity(output_iamf_)); - - RETURN_IF_NOT_OK(ObuSequencerBase::WriteTemporalUnit( - include_temporal_delimiters_, temporal_unit.second, wb, num_samples)); + const auto write_status = WriteIaSequenceToFile( + iamf_filename_, include_temporal_delimiters_, temporal_unit_map, wb); + if (!write_status.ok()) { + MaybeRemoveFile(iamf_filename_); } - LOG(INFO) << "Wrote " << temporal_unit_map.size() - << " temporal units with a total of " << num_samples - << " samples excluding padding."; - - // Flush any unwritten bytes before exiting. - RETURN_IF_NOT_OK(wb.FlushAndWriteToFile(output_iamf_)); - - return absl::OkStatus(); + return write_status; } } // namespace iamf_tools diff --git a/iamf/cli/obu_sequencer.h b/iamf/cli/obu_sequencer.h index ec1d812..13fc16d 100644 --- a/iamf/cli/obu_sequencer.h +++ b/iamf/cli/obu_sequencer.h @@ -13,9 +13,7 @@ #define CLI_OBU_SEQUENCER_H_ #include -#include #include -#include #include #include @@ -150,11 +148,7 @@ class ObuSequencerIamf : public ObuSequencerBase { bool include_temporal_delimiters, const LebGenerator& leb_generator) : ObuSequencerBase(leb_generator), - output_iamf_( - iamf_filename.empty() - ? std::nullopt - : std::make_optional( - iamf_filename, std::fstream::out | std::fstream::binary)), + iamf_filename_(iamf_filename), include_temporal_delimiters_(include_temporal_delimiters) {} ~ObuSequencerIamf() override = default; @@ -180,7 +174,7 @@ class ObuSequencerIamf : public ObuSequencerBase { const std::list& arbitrary_obus) override; private: - std::optional output_iamf_; + const std::string iamf_filename_; const bool include_temporal_delimiters_; }; diff --git a/iamf/cli/tests/obu_sequencer_test.cc b/iamf/cli/tests/obu_sequencer_test.cc index 49cb3c5..d8d912a 100644 --- a/iamf/cli/tests/obu_sequencer_test.cc +++ b/iamf/cli/tests/obu_sequencer_test.cc @@ -51,7 +51,7 @@ namespace { using ::absl_testing::IsOk; constexpr DecodedUleb128 kCodecConfigId = 1; -const uint32_t kSampleRate = 48000; +constexpr uint32_t kSampleRate = 48000; constexpr DecodedUleb128 kFirstAudioElementId = 1; constexpr DecodedUleb128 kSecondAudioElementId = 2; constexpr DecodedUleb128 kFirstSubstreamId = 1; @@ -59,7 +59,7 @@ constexpr DecodedUleb128 kSecondSubstreamId = 2; constexpr DecodedUleb128 kFirstMixPresentationId = 100; constexpr DecodedUleb128 kFirstDemixingParameterId = 998; constexpr DecodedUleb128 kCommonMixGainParameterId = 999; -const uint32_t kCommonMixGainParameterRate = kSampleRate; +constexpr uint32_t kCommonMixGainParameterRate = kSampleRate; constexpr absl::string_view kOmitOutputIamfFile = ""; constexpr bool kIncludeTemporalDelimiters = true; @@ -446,6 +446,21 @@ class ObuSequencerTest : public ::testing::Test { ASSERT_FALSE(mix_presentation_obus_.empty()); } + void InitObusForOneFrameIaSequence() { + ia_sequence_header_obu_.emplace(ObuHeader(), IASequenceHeaderObu::kIaCode, + ProfileVersion::kIamfSimpleProfile, + ProfileVersion::kIamfSimpleProfile); + per_id_metadata_ = + CreatePerIdMetadataForDemixing(kFirstDemixingParameterId); + InitializeOneParameterBlockAndOneAudioFrame( + per_id_metadata_, parameter_blocks_, audio_frames_, codec_config_obus_, + audio_elements_); + AddMixPresentationObuWithAudioElementIds( + kFirstMixPresentationId, {audio_elements_.begin()->first}, + kCommonMixGainParameterId, kCommonMixGainParameterRate, + mix_presentation_obus_); + } + void ValidateWriteDescriptorObuSequence( const std::list& expected_sequence) { WriteBitBuffer expected_wb(128); @@ -469,6 +484,11 @@ class ObuSequencerTest : public ::testing::Test { absl::flat_hash_map codec_config_obus_; absl::flat_hash_map audio_elements_; std::list mix_presentation_obus_; + + PerIdParameterMetadata per_id_metadata_; + std::list parameter_blocks_; + std::list audio_frames_; + std::list arbitrary_obus_; }; @@ -739,5 +759,85 @@ TEST(ObuSequencerIamf, PickAndPlaceSucceedsWithEmptyOutputFile) { IsOk()); } +TEST_F(ObuSequencerTest, PickAndPlaceCreatesFileWithOneFrameIaSequence) { + const std::string kOutputIamfFilename = GetAndCleanupOutputFileName(".iamf"); + InitObusForOneFrameIaSequence(); + ObuSequencerIamf sequencer(kOutputIamfFilename, + kDoNotIncludeTemporalDelimiters, + *LebGenerator::Create()); + + ASSERT_THAT( + sequencer.PickAndPlace(*ia_sequence_header_obu_, codec_config_obus_, + audio_elements_, mix_presentation_obus_, + audio_frames_, parameter_blocks_, arbitrary_obus_), + IsOk()); + + EXPECT_TRUE(std::filesystem::exists(kOutputIamfFilename)); +} + +TEST_F(ObuSequencerTest, PickAndPlaceLeavesNoFileWhenDescriptorsAreInvalid) { + constexpr uint32_t kInvalidIaCode = IASequenceHeaderObu::kIaCode + 1; + const std::string kOutputIamfFilename = GetAndCleanupOutputFileName(".iamf"); + InitObusForOneFrameIaSequence(); + // Overwrite the IA Sequence Header with an invalid one. + ia_sequence_header_obu_ = IASequenceHeaderObu( + ObuHeader(), kInvalidIaCode, ProfileVersion::kIamfSimpleProfile, + ProfileVersion::kIamfSimpleProfile); + ObuSequencerIamf sequencer(kOutputIamfFilename, + kDoNotIncludeTemporalDelimiters, + *LebGenerator::Create()); + + ASSERT_FALSE(sequencer + .PickAndPlace(*ia_sequence_header_obu_, codec_config_obus_, + audio_elements_, mix_presentation_obus_, + audio_frames_, parameter_blocks_, + arbitrary_obus_) + .ok()); + + EXPECT_FALSE(std::filesystem::exists(kOutputIamfFilename)); +} + +TEST_F(ObuSequencerTest, PickAndPlaceLeavesNoFileWhenTemporalUnitsAreInvalid) { + constexpr bool kInvalidateTemporalUnit = true; + const std::string kOutputIamfFilename = GetAndCleanupOutputFileName(".iamf"); + InitObusForOneFrameIaSequence(); + arbitrary_obus_.emplace_back( + ArbitraryObu(kObuIaReserved25, ObuHeader(), {}, + ArbitraryObu::kInsertionHookAfterAudioFramesAtTick, 0, + kInvalidateTemporalUnit)); + ObuSequencerIamf sequencer(kOutputIamfFilename, + kDoNotIncludeTemporalDelimiters, + *LebGenerator::Create()); + + ASSERT_FALSE(sequencer + .PickAndPlace(*ia_sequence_header_obu_, codec_config_obus_, + audio_elements_, mix_presentation_obus_, + audio_frames_, parameter_blocks_, + arbitrary_obus_) + .ok()); + + EXPECT_FALSE(std::filesystem::exists(kOutputIamfFilename)); +} + +TEST_F(ObuSequencerTest, + PickAndPlaceOnInvalidTemporalUnitFailsWhenOutputFileIsOmitted) { + constexpr bool kInvalidateTemporalUnit = true; + InitObusForOneFrameIaSequence(); + arbitrary_obus_.emplace_back( + ArbitraryObu(kObuIaReserved25, ObuHeader(), {}, + ArbitraryObu::kInsertionHookAfterAudioFramesAtTick, 0, + kInvalidateTemporalUnit)); + ObuSequencerIamf sequencer(std::string(kOmitOutputIamfFile), + kDoNotIncludeTemporalDelimiters, + *LebGenerator::Create()); + + ASSERT_FALSE(sequencer + .PickAndPlace(*ia_sequence_header_obu_, codec_config_obus_, + audio_elements_, mix_presentation_obus_, + audio_frames_, parameter_blocks_, + arbitrary_obus_) + .ok()); +} + } // namespace } // namespace iamf_tools