Extract bit depth from MP4s

src/bit_reader.rs

@@ -24,6 +24,20 @@ impl<'a> BitReader<'a> {
   pub(crate) fn new(bytes: &'a [u8]) -> Self {
     Self { bytes, i: 0 }
   }
+
+  pub(crate) fn ue(&mut self) -> Option<u64> {
+    let mut zeros = 0;
+
+    while self.bit()? == 0 {
+      zeros += 1;
+
+      if zeros > 31 {
+        return None;
+      }
+    }
+
+    Some((1 << zeros) - 1 + self.bits(zeros)?)
+  }
 }
 
 #[cfg(test)]
@@ -39,4 +53,19 @@ mod tests {
     assert_eq!(reader.bits(6), Some(0b01_1011));
     assert_eq!(reader.bits(7), None);
   }
+
+  #[test]
+  fn ue() {
+    #[track_caller]
+    fn case(data: &[u8], expected: Option<u64>) {
+      assert_eq!(BitReader::new(data).ue(), expected);
+    }
+
+    case(&[0b1000_0000], Some(0));
+    case(&[0b0100_0000], Some(1));
+    case(&[0b0110_0000], Some(2));
+    case(&[0b0001_0010], Some(8));
+    case(&[0b0000_0000], None);
+    case(&[], None);
+  }
 }

src/mp4_builder.rs

@@ -1,8 +1,10 @@
 pub struct Mp4Builder {
+  avcc_profile: u8,
   duration: u32,
   frame_count: u32,
   sample_size: u32,
   sample_sizes: Vec<u32>,
+  sps: Vec<u8>,
   timescale: u32,
   tracks: Vec<Vec<u8>>,
 }
@@ -45,6 +47,12 @@ impl Mp4Builder {
     self.track(*b"soun", &[entry])
   }
 
+  #[must_use]
+  pub fn avcc_profile(mut self, avcc_profile: u8) -> Self {
+    self.avcc_profile = avcc_profile;
+    self
+  }
+
   pub fn build(self) -> Vec<u8> {
     let mut ftyp = Vec::new();
     ftyp.extend_from_slice(b"isom");
@@ -76,10 +84,12 @@ impl Mp4Builder {
 
   pub fn new() -> Self {
     Self {
+      avcc_profile: 0,
       duration: 0,
       frame_count: 0,
       sample_size: 1,
       sample_sizes: Vec::new(),
+      sps: Vec::new(),
       timescale: 1000,
       tracks: Vec::new(),
     }
@@ -98,6 +108,12 @@ impl Mp4Builder {
     self
   }
 
+  #[must_use]
+  pub fn sps(mut self, sps: &[u8]) -> Self {
+    self.sps = sps.into();
+    self
+  }
+
   #[must_use]
   pub fn timescale(mut self, timescale: u32) -> Self {
     self.timescale = timescale;
@@ -190,22 +206,38 @@ impl Mp4Builder {
     self
   }
 
-  pub fn video_entry(entry: [u8; 4], config: [u8; 4], width: u16, height: u16) -> Vec<u8> {
+  pub fn video_entry(
+    entry: [u8; 4],
+    config: [u8; 4],
+    config_payload: &[u8],
+    width: u16,
+    height: u16,
+  ) -> Vec<u8> {
     let mut payload = Vec::new();
     payload.extend_from_slice(&[0; 6]);
     payload.extend_from_slice(&[0, 1]);
     payload.extend_from_slice(&[0; 16]);
     payload.extend_from_slice(&width.to_be_bytes());
     payload.extend_from_slice(&height.to_be_bytes());
     payload.extend_from_slice(&[0; 50]);
-    payload.extend_from_slice(&Self::atom(config, &[1, 0, 0, 0, 0xff, 0xe0, 0]));
+    payload.extend_from_slice(&Self::atom(config, config_payload));
 
     Self::atom(entry, &payload)
   }
 
   #[must_use]
   pub fn video_track(self, width: u16, height: u16) -> Self {
-    let entry = Self::video_entry(*b"avc1", *b"avcC", width, height);
+    let avcc = if self.sps.is_empty() {
+      vec![1, self.avcc_profile, 0, 0, 0xff, 0xe0, 0]
+    } else {
+      let mut avcc = vec![1, self.sps[1], self.sps[2], self.sps[3], 0xff, 0xe1];
+      avcc.extend_from_slice(&u16::try_from(self.sps.len()).unwrap().to_be_bytes());
+      avcc.extend_from_slice(&self.sps);
+      avcc.push(0);
+      avcc
+    };
+
+    let entry = Self::video_entry(*b"avc1", *b"avcC", &avcc, width, height);
     self.track(*b"vide", &[entry])
   }
 }

src/video.rs

@@ -46,6 +46,54 @@ impl Video {
     formats
   }
 
+  fn h264_bit_depth(sps: &[u8]) -> Option<u64> {
+    let mut rbsp = Vec::new();
+
+    // skip NAL unit header
+    for &byte in sps.get(1..)? {
+      // remove emulation prevention bytes
+      if byte == 3 && rbsp.ends_with(&[0, 0]) {
+        continue;
+      }
+
+      rbsp.push(byte);
+    }
+
+    let mut reader = BitReader::new(&rbsp);
+
+    // profile_idc
+    let profile_idc = reader.bits(8)?;
+
+    // constraint flags
+    reader.bits(8)?;
+
+    // level_idc
+    reader.bits(8)?;
+
+    // seq_parameter_set_id
+    reader.ue()?;
+
+    if !Self::h264_high_profile(profile_idc) {
+      return Some(8);
+    }
+
+    // chroma_format_idc
+    if reader.ue()? == 3 {
+      // separate_colour_plane_flag
+      reader.bit()?;
+    }
+
+    // bit_depth_luma_minus8
+    Some(8 + reader.ue()?)
+  }
+
+  fn h264_high_profile(profile_idc: u64) -> bool {
+    matches!(
+      profile_idc,
+      44 | 83 | 86 | 100 | 110 | 118 | 122 | 128 | 134 | 135 | 138 | 139 | 244
+    )
+  }
+
   fn info(&self, root: &Utf8Path) -> Result<VideoInfo> {
     let path = root.join(self.as_path());
 
@@ -163,10 +211,21 @@ impl Video {
             );
           };
 
+          let bit_depth = if let Some(sps) = avc1.avcc.sequence_parameter_sets.first() {
+            Some(Self::h264_bit_depth(&sps.bytes).context(video_error::SpsInvalid)?)
+          } else {
+            ensure!(
+              !Self::h264_high_profile(avc1.avcc.avc_profile_indication.into()),
+              video_error::SpsMissing,
+            );
+
+            Some(8)
+          };
+
           video_track = Some(Track {
             codec: Codec::H264,
             info: TrackInfo::Video {
-              bit_depth: None,
+              bit_depth,
               dimensions: Dimensions {
                 height: avc1.height.into(),
                 width: avc1.width.into(),
@@ -641,6 +700,22 @@ mod tests {
     case("foo/bar.mp4", ComponentError::Separator { character: '/' });
   }
 
+  #[test]
+  fn h264_bit_depth() {
+    #[track_caller]
+    fn case(sps: &[u8], expected: Option<u64>) {
+      assert_eq!(Video::h264_bit_depth(sps), expected);
+    }
+
+    case(&[0x67, 66, 0, 30, 0x80], Some(8));
+    case(&[0x67, 100, 0, 31, 0xA6], Some(10));
+    case(&[0x67, 100, 0, 31, 0x91], Some(8));
+    case(&[0x67, 100, 0, 0, 0x03, 0xA6], Some(10));
+    case(&[0x67, 100, 0, 31], None);
+    case(&[0x67], None);
+    case(&[], None);
+  }
+
   #[test]
   fn mp4_info() {
     #[track_caller]
@@ -663,7 +738,7 @@ mod tests {
           Track {
             codec: Codec::H264,
             info: TrackInfo::Video {
-              bit_depth: None,
+              bit_depth: Some(8),
               dimensions: Dimensions {
                 height: 1,
                 width: 2,
@@ -688,7 +763,7 @@ mod tests {
         tracks: vec![Track {
           codec: Codec::H264,
           info: TrackInfo::Video {
-            bit_depth: None,
+            bit_depth: Some(8),
             dimensions: Dimensions {
               height: 1,
               width: 2,
@@ -725,7 +800,7 @@ mod tests {
         .tracks[0]
         .info,
       TrackInfo::Video {
-        bit_depth: None,
+        bit_depth: Some(8),
         dimensions: Dimensions {
           height: 1,
           width: 2,
@@ -754,7 +829,7 @@ mod tests {
       Track {
         codec: Codec::H264,
         info: TrackInfo::Video {
-          bit_depth: None,
+          bit_depth: Some(8),
           dimensions: Dimensions {
             height: 1,
             width: 2,
@@ -765,6 +840,35 @@ mod tests {
       },
     );
 
+    assert_eq!(
+      case(
+        Mp4Builder::new()
+          .sps(&[0x67, 100, 0, 31, 0xA6])
+          .video_track(2, 1),
+      )
+      .unwrap()
+      .tracks[0]
+        .info,
+      TrackInfo::Video {
+        bit_depth: Some(10),
+        dimensions: Dimensions {
+          height: 1,
+          width: 2,
+        },
+        frames: 0,
+      },
+    );
+
+    error(
+      Mp4Builder::new().sps(&[0x67, 100, 0, 31]).video_track(2, 1),
+      "invalid SPS",
+    );
+
+    error(
+      Mp4Builder::new().avcc_profile(100).video_track(2, 1),
+      "missing SPS",
+    );
+
     error(
       Mp4Builder::new().timescale(0).video_track(2, 1),
       "zero timescale",
@@ -796,7 +900,13 @@ mod tests {
       Mp4Builder::new()
         .track(
           *b"vide",
-          &[Mp4Builder::video_entry(*b"s263", *b"d263", 2, 1)],
+          &[Mp4Builder::video_entry(
+            *b"s263",
+            *b"d263",
+            &[1, 0, 0, 0, 0xff, 0xe0, 0],
+            2,
+            1,
+          )],
         )
         .audio_track(0x40),
       "track 0 has unsupported video codec `s263`",
@@ -843,7 +953,7 @@ mod tests {
           Track {
             codec: Codec::H264,
             info: TrackInfo::Video {
-              bit_depth: None,
+              bit_depth: Some(8),
               dimensions: Dimensions {
                 height: 1,
                 width: 2,
@@ -875,7 +985,7 @@ mod tests {
         Track {
           codec: Codec::H264,
           info: TrackInfo::Video {
-            bit_depth: None,
+            bit_depth: Some(8),
             dimensions: Dimensions {
               height: 1,
               width: 2,

src/video_error.rs

@@ -21,6 +21,10 @@ pub enum VideoError {
   DurationMissing,
   #[snafu(display("duration overflow"))]
   DurationOverflow,
+  #[snafu(display("invalid SPS"))]
+  SpsInvalid,
+  #[snafu(display("missing SPS"))]
+  SpsMissing,
   #[snafu(display("zero timescale"))]
   TimescaleZero,
   #[snafu(display("unsupported timestamp scale {timestamp_scale}"))]

tests/metadata.rs

@@ -171,6 +171,7 @@ media:
             "codec": "h264",
             "info": {
               "type": "video",
+              "bit_depth": 8,
               "dimensions": {
                 "height": 1,
                 "width": 2