feat: Write colr atom to muxed mp4 (#1261)

This PR is an extension of the full AV1 codec string feature: [PR
1205](https://github.com/shaka-project/shaka-packager/pull/1205) and
relates to [Issue
1007](https://github.com/shaka-project/shaka-packager/issues/1007) and
[Issue
1202](https://github.com/shaka-project/shaka-packager/issues/1202).

As per the AV1 spec, the codec string may contain optional color values.
These color values are critical for detecting HDR video streams - see
[Issue
1007](https://github.com/shaka-project/shaka-packager/issues/1007).
Color information is extracted from the input mp4's `colr` atom and used
to generate the full AV1 codec string. This PR preserves the color
information by writing the `colr` atom to the muxed mp4.

**References**:
- [AV1 Codec ISO Media File Format
Binding](https://aomediacodec.github.io/av1-isobmff/#codecsparam)
- [AV1 Bitstream & Decoding Process
Specification - Section 6.4.2 Color config semantics (page
117)](https://aomediacodec.github.io/av1-spec/av1-spec.pdf)
- [QuickTime File Format
Specification](https://developer.apple.com/library/archive/documentation/QuickTime/QTFF/QTFFChap3/qtff3.html#//apple_ref/doc/uid/TP40000939-CH205-125526)

packager/media/base/video_stream_info.h

@@ -69,6 +69,7 @@ class VideoStreamInfo : public StreamInfo {
   uint32_t trick_play_factor() const { return trick_play_factor_; }
   uint32_t playback_rate() const { return playback_rate_; }
   const std::vector<uint8_t>& eme_init_data() const { return eme_init_data_; }
+  const std::vector<uint8_t>& colr_data() const { return colr_data_; }
 
   void set_extra_config(const std::vector<uint8_t>& extra_config) {
     extra_config_ = extra_config;
@@ -90,6 +91,9 @@ class VideoStreamInfo : public StreamInfo {
                          size_t eme_init_data_size) {
     eme_init_data_.assign(eme_init_data, eme_init_data + eme_init_data_size);
   }
+  void set_colr_data(const uint8_t* colr_data, size_t colr_data_size) {
+    colr_data_.assign(colr_data, colr_data + colr_data_size);
+  }
 
  private:
   // Extra codec configuration in a stream of mp4 boxes. It is only applicable
@@ -128,6 +132,9 @@ class VideoStreamInfo : public StreamInfo {
   // https://w3c.github.io/encrypted-media/#initialization-data.
   std::vector<uint8_t> eme_init_data_;
 
+  // Raw colr atom data. It is only applicable to the mp4 container.
+  std::vector<uint8_t> colr_data_;
+
   // Not using DISALLOW_COPY_AND_ASSIGN here intentionally to allow the compiler
   // generated copy constructor and assignment operator. Since the extra data is
   // typically small, the performance impact is minimal.

packager/media/formats/mp4/box_definitions.cc

@@ -1470,29 +1470,33 @@ FourCC ColorParameters::BoxType() const {
 }
 
 bool ColorParameters::ReadWriteInternal(BoxBuffer* buffer) {
-  if (buffer->reader()) {
+  if (buffer->Reading()) {
-    RCHECK((buffer->reader())->ReadFourCC(&color_parameter_type) &&
+    BoxReader* reader = buffer->reader();
-           (buffer->reader())->Read2(&color_primaries) &&
+    DCHECK(reader);
-           (buffer->reader())->Read2(&transfer_characteristics) &&
+
-           (buffer->reader())->Read2(&matrix_coefficients));
+    // Parse and store the raw box for colr atom preservation in the output mp4.
+    raw_box.assign(reader->data(), reader->data() + reader->size());
+
+    // Parse individual parameters for full codec string formation.
+    RCHECK(reader->ReadFourCC(&color_parameter_type) &&
+           reader->Read2(&color_primaries) &&
+           reader->Read2(&transfer_characteristics) &&
+           reader->Read2(&matrix_coefficients));
     // Type nclc does not contain video_full_range_flag data, and thus, it has 1
     // less byte than nclx. Only extract video_full_range_flag if of type nclx.
     if (color_parameter_type == FOURCC_nclx) {
-      RCHECK((buffer->reader())->Read1(&video_full_range_flag));
+      RCHECK(reader->Read1(&video_full_range_flag));
     }
+  } else {
+    // When writing, only need to write the raw_box.
+    DCHECK(!raw_box.empty());
+    buffer->writer()->AppendVector(raw_box);
   }
-  // TODO(caitlinocallaghan) Add the ability to write the colr atom and include
-  // it in the muxed mp4.
   return true;
 }
 
 size_t ColorParameters::ComputeSizeInternal() {
-  // This box is optional. Skip it if it is not initialized.
+  return raw_box.size();
-  if (color_parameter_type == FOURCC_NULL)
-    return 0;
-  return HeaderSize() + kFourCCSize + sizeof(color_primaries) +
-         sizeof(transfer_characteristics) + sizeof(matrix_coefficients) +
-         sizeof(video_full_range_flag);
 }
 
 PixelAspectRatio::PixelAspectRatio() = default;
@@ -1652,9 +1656,10 @@ size_t VideoSampleEntry::ComputeSizeInternal() {
   size_t size = HeaderSize() + sizeof(data_reference_index) + sizeof(width) +
                 sizeof(height) + sizeof(kVideoResolution) * 2 +
                 sizeof(kVideoFrameCount) + sizeof(kVideoDepth) +
-                pixel_aspect.ComputeSize() + sinf.ComputeSize() +
+                colr.ComputeSize() + pixel_aspect.ComputeSize() +
-                codec_configuration.ComputeSize() + kCompressorNameSize + 6 +
+                sinf.ComputeSize() + codec_configuration.ComputeSize() +
-                4 + 16 + 2;  // 6 + 4 bytes reserved, 16 + 2 bytes predefined.
+                kCompressorNameSize + 6 + 4 + 16 +
+                2;  // 6 + 4 bytes reserved, 16 + 2 bytes predefined.
   for (CodecConfiguration& codec_config : extra_codec_configs)
     size += codec_config.ComputeSize();
   return size;

packager/media/formats/mp4/box_definitions.h

@@ -276,6 +276,7 @@ struct ColorParameters : Box {
   uint16_t transfer_characteristics = 1;
   uint16_t matrix_coefficients = 1;
   uint8_t video_full_range_flag = 0;
+  std::vector<uint8_t> raw_box;
 };
 
 struct PixelAspectRatio : Box {

packager/media/formats/mp4/mp4_media_parser.cc

@@ -726,6 +726,8 @@ bool MP4MediaParser::ParseMoov(BoxReader* reader) {
           0,  // trick_play_factor
           nalu_length_size, track->media.header.language.code, is_encrypted));
       video_stream_info->set_extra_config(entry.ExtraCodecConfigsAsVector());
+      video_stream_info->set_colr_data((entry.colr.raw_box).data(),
+                                       (entry.colr.raw_box).size());
 
       // Set pssh raw data if it has.
       if (moov_->pssh.size() > 0) {

packager/media/formats/mp4/mp4_muxer.cc

@@ -431,6 +431,7 @@ bool MP4Muxer::GenerateVideoTrak(const VideoStreamInfo* video_info,
       CodecToFourCC(video_info->codec(), video_info->h26x_stream_format());
   video.width = video_info->width();
   video.height = video_info->height();
+  video.colr.raw_box = video_info->colr_data();
   video.codec_configuration.data = video_info->codec_config();
   if (!video.ParseExtraCodecConfigsVector(video_info->extra_config())) {
     LOG(ERROR) << "Malformed extra codec configs: "