Add logic to decode temporal units in IamfDecoder::Decode. This imple…

…mentation handles streams in which temporal units are always whole; it notably does not handle cases in which Decode() receives a partial temporal unit. This case is complicated and we currently don't have the proper underlying code to support this at this level. Subsequent cls will fix this issue.

PiperOrigin-RevId: 728761230

iamf/api/BUILD

@@ -8,7 +8,9 @@ cc_library(
     srcs = ["iamf_decoder.cc"],
     hdrs = ["iamf_decoder.h"],
     deps = [
+        "//iamf/cli:audio_frame_with_data",
         "//iamf/cli:obu_processor",
+        "//iamf/cli:parameter_block_with_data",
         "//iamf/cli:rendering_mix_presentation_finalizer",
         "//iamf/common:read_bit_buffer",
         "//iamf/common/utils:macros",

iamf/api/iamf_decoder.cc

@@ -16,13 +16,16 @@
 #include <iterator>
 #include <list>
 #include <memory>
+#include <optional>
 #include <utility>
 #include <vector>
 
 #include "absl/log/log.h"
 #include "absl/status/status.h"
 #include "absl/types/span.h"
+#include "iamf/cli/audio_frame_with_data.h"
 #include "iamf/cli/obu_processor.h"
+#include "iamf/cli/parameter_block_with_data.h"
 #include "iamf/cli/rendering_mix_presentation_finalizer.h"
 #include "iamf/common/read_bit_buffer.h"
 #include "iamf/common/utils/macros.h"
@@ -65,6 +68,46 @@ absl::StatusOr<std::unique_ptr<ObuProcessor>> CreateObuProcessor(
   return obu_processor;
 }
 
+absl::Status ProcessAllTemporalUnits(
+    StreamBasedReadBitBuffer* read_bit_buffer, ObuProcessor* obu_processor,
+    std::vector<std::vector<std::vector<int32_t>>>& rendered_pcm_samples) {
+  LOG(INFO) << "Processing Temporal Units";
+  int32_t num_bits_read = 0;
+  bool continue_processing = true;
+  while (continue_processing) {
+    auto start_position_for_temporal_unit = read_bit_buffer->Tell();
+    std::list<AudioFrameWithData> audio_frames_for_temporal_unit;
+    std::list<ParameterBlockWithData> parameter_blocks_for_temporal_unit;
+    std::optional<int32_t> timestamp_for_temporal_unit;
+    // TODO(b/395889878): Add support for partial temporal units.
+    RETURN_IF_NOT_OK(obu_processor->ProcessTemporalUnit(
+        audio_frames_for_temporal_unit, parameter_blocks_for_temporal_unit,
+        timestamp_for_temporal_unit, continue_processing));
+
+    // Trivial IA Sequences may have empty temporal units. Do not try to
+    // render empty temporal unit.
+    if (timestamp_for_temporal_unit.has_value()) {
+      absl::Span<const std::vector<int32_t>>
+          rendered_pcm_samples_for_temporal_unit;
+      RETURN_IF_NOT_OK(obu_processor->RenderTemporalUnitAndMeasureLoudness(
+          *timestamp_for_temporal_unit, audio_frames_for_temporal_unit,
+          parameter_blocks_for_temporal_unit,
+          rendered_pcm_samples_for_temporal_unit));
+      rendered_pcm_samples.push_back(
+          std::vector(rendered_pcm_samples_for_temporal_unit.begin(),
+                      rendered_pcm_samples_for_temporal_unit.end()));
+    }
+    num_bits_read +=
+        (read_bit_buffer->Tell() - start_position_for_temporal_unit);
+  }
+  // Empty the buffer of the data that was processed thus far.
+  RETURN_IF_NOT_OK(read_bit_buffer->Flush(num_bits_read / 8));
+  LOG(INFO) << "Rendered " << rendered_pcm_samples.size()
+            << " temporal units. Please call GetOutputTemporalUnit() to get "
+               "the rendered PCM samples.";
+  return absl::OkStatus();
+}
+
 }  // namespace
 
 absl::StatusOr<IamfDecoder> IamfDecoder::Create() {
@@ -111,6 +154,10 @@ absl::Status IamfDecoder::Decode(absl::Span<const uint8_t> bitstream) {
       return obu_processor.status();
     }
   }
+
+  // At this stage, we know that we've processed all descriptor OBUs.
+  RETURN_IF_NOT_OK(ProcessAllTemporalUnits(
+      read_bit_buffer_.get(), obu_processor_.get(), rendered_pcm_samples_));
   return absl::OkStatus();
 }
 

iamf/api/iamf_decoder.h

@@ -251,6 +251,12 @@ class IamfDecoder {
 
   // Buffer that is filled with data from Decode().
   std::unique_ptr<StreamBasedReadBitBuffer> read_bit_buffer_;
+
+  // Rendered PCM samples. Each element in the outer vector corresponds to a
+  // temporal unit. A temporal unit will never be partially filled, so the
+  // number of elements in the outer vector is equal to the number of decoded
+  // temporal units currently available.
+  std::vector<std::vector<std::vector<int32_t>>> rendered_pcm_samples_;
 };
 }  // namespace iamf_tools
 

iamf/api/tests/BUILD

@@ -7,6 +7,7 @@ cc_test(
         "//iamf/cli:audio_element_with_data",
         "//iamf/cli/tests:cli_test_utils",
         "//iamf/common:write_bit_buffer",
+        "//iamf/obu:audio_frame",
         "//iamf/obu:codec_config",
         "//iamf/obu:ia_sequence_header",
         "//iamf/obu:mix_presentation",

iamf/api/tests/iamf_decoder_test.cc

@@ -12,6 +12,7 @@
 
 #include "iamf/api/iamf_decoder.h"
 
+#include <array>
 #include <cstdint>
 #include <list>
 #include <vector>
@@ -24,6 +25,7 @@
 #include "iamf/cli/audio_element_with_data.h"
 #include "iamf/cli/tests/cli_test_utils.h"
 #include "iamf/common/write_bit_buffer.h"
+#include "iamf/obu/audio_frame.h"
 #include "iamf/obu/codec_config.h"
 #include "iamf/obu/ia_sequence_header.h"
 #include "iamf/obu/mix_presentation.h"
@@ -44,6 +46,8 @@ constexpr DecodedUleb128 kFirstSubstreamId = 18;
 constexpr DecodedUleb128 kFirstMixPresentationId = 3;
 constexpr DecodedUleb128 kCommonMixGainParameterId = 999;
 constexpr DecodedUleb128 kCommonParameterRate = kSampleRate;
+constexpr std::array<uint8_t, 16> kEightSampleAudioFrame = {
+    1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16};
 
 // TODO(b/396453922): Move this to a common test utils file.
 std::vector<uint8_t> SerializeObus(
@@ -62,7 +66,8 @@ std::vector<uint8_t> GenerateBasicDescriptorObus() {
       ObuHeader(), IASequenceHeaderObu::kIaCode,
       ProfileVersion::kIamfSimpleProfile, ProfileVersion::kIamfBaseProfile);
   absl::flat_hash_map<DecodedUleb128, CodecConfigObu> codec_configs;
-  AddOpusCodecConfigWithId(kFirstCodecConfigId, codec_configs);
+  AddLpcmCodecConfigWithIdAndSampleRate(kFirstCodecConfigId, kSampleRate,
+                                        codec_configs);
   absl::flat_hash_map<DecodedUleb128, AudioElementWithData> audio_elements;
   AddAmbisonicsMonoAudioElementWithSubstreamIds(
       kFirstAudioElementId, kFirstCodecConfigId, {kFirstSubstreamId},
@@ -99,6 +104,14 @@ TEST(CreateFromDescriptors, Succeeds) {
   EXPECT_TRUE(decoder->IsDescriptorProcessingComplete());
 }
 
+TEST(CreateFromDescriptors, FailsWithIncompleteDescriptorObus) {
+  auto descriptors = GenerateBasicDescriptorObus();
+  // remove the last byte to make the descriptor OBUs incomplete.
+  descriptors.pop_back();
+  auto decoder = IamfDecoder::CreateFromDescriptors(descriptors);
+  EXPECT_FALSE(decoder.ok());
+}
+
 TEST(Decode, SucceedsAndProcessesDescriptorsWithTemporalDelimiterAtEnd) {
   auto decoder = IamfDecoder::Create();
   ASSERT_THAT(decoder, IsOk());
@@ -132,9 +145,7 @@ TEST(Decode, SucceedsWithMultiplePushesOfDescriptorObus) {
   EXPECT_TRUE(decoder->IsDescriptorProcessingComplete());
 }
 
-// TODO(b/396466476): Switch to failing this test once we process temporal units
-// and detect extra descriptor OBUs in a contract-breaking way.
-TEST(CreateFromDescriptors, SucceedsWithDescriptorObuInSubsequentDecode) {
+TEST(CreateFromDescriptors, FailsWithDescriptorObuInSubsequentDecode) {
   auto decoder =
       IamfDecoder::CreateFromDescriptors(GenerateBasicDescriptorObus());
   EXPECT_THAT(decoder, IsOk());
@@ -145,15 +156,45 @@ TEST(CreateFromDescriptors, SucceedsWithDescriptorObuInSubsequentDecode) {
       kFirstMixPresentationId + 1, {kFirstAudioElementId},
       kCommonMixGainParameterId, kCommonParameterRate, mix_presentation_obus);
   auto second_chunk = SerializeObus({&mix_presentation_obus.front()});
-  EXPECT_THAT(decoder->Decode(second_chunk), IsOk());
+
+  EXPECT_FALSE(decoder->Decode(second_chunk).ok());
 }
 
-TEST(CreateFromDescriptors, FailsWithIncompleteDescriptorObus) {
-  auto descriptors = GenerateBasicDescriptorObus();
-  // remove the last byte to make the descriptor OBUs incomplete.
-  descriptors.pop_back();
-  auto decoder = IamfDecoder::CreateFromDescriptors(descriptors);
-  EXPECT_FALSE(decoder.ok());
+TEST(Decode, SucceeedsWithSeparatePushesOfDescriptorAndTemporalUnits) {
+  std::vector<uint8_t> source_data = GenerateBasicDescriptorObus();
+  auto decoder = IamfDecoder::CreateFromDescriptors(source_data);
+  ASSERT_THAT(decoder, IsOk());
+  AudioFrameObu audio_frame(ObuHeader(), kFirstSubstreamId,
+                            kEightSampleAudioFrame);
+  auto temporal_unit = SerializeObus({&audio_frame});
+
+  EXPECT_THAT(decoder->Decode(temporal_unit), IsOk());
+}
+
+TEST(Decode, SucceedsWithOneTemporalUnit) {
+  auto decoder = IamfDecoder::Create();
+  ASSERT_THAT(decoder, IsOk());
+  std::vector<uint8_t> source_data = GenerateBasicDescriptorObus();
+  AudioFrameObu audio_frame(ObuHeader(), kFirstSubstreamId,
+                            kEightSampleAudioFrame);
+  auto temporal_unit = SerializeObus({&audio_frame});
+  source_data.insert(source_data.end(), temporal_unit.begin(),
+                     temporal_unit.end());
+
+  EXPECT_THAT(decoder->Decode(source_data), IsOk());
+}
+
+TEST(Decode, SucceedsWithMultipleTemporalUnits) {
+  auto decoder = IamfDecoder::Create();
+  ASSERT_THAT(decoder, IsOk());
+  std::vector<uint8_t> source_data = GenerateBasicDescriptorObus();
+  AudioFrameObu audio_frame(ObuHeader(), kFirstSubstreamId,
+                            kEightSampleAudioFrame);
+  auto temporal_units = SerializeObus({&audio_frame, &audio_frame});
+  source_data.insert(source_data.end(), temporal_units.begin(),
+                     temporal_units.end());
+
+  EXPECT_THAT(decoder->Decode(source_data), IsOk());
 }
 
 }  // namespace