progress on parsing P macroblocks

Author: Djuffin

src/h264/decoder.rs

@@ -107,10 +107,6 @@ impl Decoder {
                             .map_err(parse_error_handler)?;
 
                     info!("non-IDR Slice: {:#?}", slice);
-                    if slice.header.slice_type != SliceType::I {
-                        todo!("implement P-slices")
-                    }
-
                     let frame = VideoFrame::new_with_padding(
                         slice.sps.pic_width(),
                         slice.sps.pic_hight(),

src/h264/macroblock.rs

@@ -378,6 +378,9 @@ pub struct IMb {
 #[derive(Clone, Debug, Default)]
 pub struct PMb {
     pub mb_type: PMbType,
+    pub coded_block_pattern: CodedBlockPattern,
+    pub mb_qp_delta: i32,
+    pub residual: Option<Box<Residual>>,
 }
 
 #[derive(Clone, Debug)]
@@ -405,6 +408,20 @@ impl IMb {
     }
 }
 
+#[allow(non_snake_case)]
+impl PMb {
+    #[inline]
+    pub fn MbPartPredMode(&self, partition: usize) -> MbPredictionMode {
+        match self.mb_type {
+            PMbType::P_Skip
+            | PMbType::P_L0_16x16
+            | PMbType::P_L0_L0_16x8
+            | PMbType::P_L0_L0_8x16 => MbPredictionMode::Pred_L0,
+            PMbType::P_8x8ref0 | PMbType::P_8x8 => MbPredictionMode::None,
+        }
+    }
+}
+
 #[allow(non_snake_case)]
 impl Macroblock {
     #[inline]
@@ -433,10 +450,8 @@ impl Macroblock {
     pub fn get_coded_block_pattern(&self) -> CodedBlockPattern {
         match self {
             Macroblock::I(mb) => mb.coded_block_pattern,
+            Macroblock::P(mb) => mb.coded_block_pattern,
             Macroblock::PCM(_) => CodedBlockPattern::default(),
-            Macroblock::P(_) => {
-                todo!("P blocks")
-            }
         }
     }
 
@@ -445,10 +460,10 @@ impl Macroblock {
             Macroblock::I(mb) => {
                 mb.residual = r;
             }
-            Macroblock::PCM(_) => {}
-            Macroblock::P(_) => {
-                todo!("P blocks")
+            Macroblock::P(mb) => {
+                mb.residual = r;
             }
+            Macroblock::PCM(_) => {}
         }
     }
 }

src/h264/parser.rs

@@ -744,9 +744,52 @@ fn calc_prev_intra4x4_pred_mode(
 }
 
 pub fn parse_macroblock(input: &mut BitReader, slice: &Slice) -> ParseResult<Macroblock> {
-    let mb_type: IMbType;
-    read_value!(input, mb_type, ue);
+    let mb_type_val: u32;
+    read_value!(input, mb_type_val, ue);
 
+    if slice.header.slice_type == SliceType::I || slice.header.slice_type == SliceType::SI {
+        let mb_type = IMbType::try_from(mb_type_val)?;
+        parse_i_macroblock(input, slice, mb_type)
+    } else {
+        let mb_type = PMbType::try_from(mb_type_val)?;
+        parse_p_macroblock(input, slice, mb_type)
+    }
+}
+
+pub fn parse_p_macroblock(
+    input: &mut BitReader,
+    slice: &Slice,
+    mb_type: PMbType,
+) -> ParseResult<Macroblock> {
+    let mut block = PMb { mb_type, ..PMb::default() };
+    let this_mb_addr = slice.get_next_mb_addr();
+    /*
+        TODO: Parse motion vectors here
+    */
+
+    let coded_block_pattern_num: u8;
+    read_value!(input, coded_block_pattern_num, ue, 8);
+    block.coded_block_pattern =
+        tables::code_num_to_inter_coded_block_pattern(coded_block_pattern_num)
+            .ok_or("Invalid coded_block_pattern")?;
+
+    if !block.coded_block_pattern.is_zero() {
+        read_value!(input, block.mb_qp_delta, se);
+        let mut residual = Box::<Residual>::default();
+        residual.coded_block_pattern = block.coded_block_pattern;
+        residual.prediction_mode = block.MbPartPredMode(0);
+        parse_residual(input, slice, &mut residual)?;
+        block.residual = Some(residual);
+    }
+
+    Ok(Macroblock::P(block))
+}
+
+pub fn parse_i_macroblock(
+    input: &mut BitReader,
+    slice: &Slice,
+    mb_type: IMbType,
+) -> ParseResult<Macroblock> {
     if mb_type == IMbType::I_PCM {
         let mut block = PcmMb::default();
         input.align();
@@ -865,7 +908,7 @@ pub fn parse_slice_data(input: &mut BitReader, slice: &mut Slice) -> ParseResult
                 ));
             }
             for i in 0..mb_skip_run {
-                let block = Macroblock::P(PMb { mb_type: PMbType::P_Skip });
+                let block = Macroblock::P(PMb { mb_type: PMbType::P_Skip, ..Default::default() });
                 slice.append_mb(block);
             }
             if mb_skip_run > 0 {

src/h264/tables.rs

@@ -22,11 +22,24 @@ const CODE_NUM_TO_INTRA_CODED_BLOCK_PATTERN: [u8; 48] = [
     35, 37, 42, 44, 1, 2, 4, 8, 17, 18, 20, 24, 6, 9, 22, 25, 32, 33, 34, 36, 40, 38, 41,
 ];
 
+const CODE_NUM_TO_INTER_CODED_BLOCK_PATTERN: [u8; 48] = [
+    0, 16, 1, 2, 4, 8, 32, 3, 5, 10, 12, 15, 47, 7, 11, 13, 14, 6, 9, 31, 35, 37, 42, 44, 33, 34,
+    36, 40, 39, 43, 45, 46, 17, 18, 20, 24, 19, 21, 26, 28, 23, 27, 29, 30, 22, 25, 38, 41,
+];
+
 #[inline]
 pub fn code_num_to_intra_coded_block_pattern(x: u8) -> Option<CodedBlockPattern> {
     let result =
         CODE_NUM_TO_INTRA_CODED_BLOCK_PATTERN.get(x as usize).map(|x| CodedBlockPattern(*x));
-    trace!("coded_block_pattern num: {} pattern: {:?}", x, result);
+    trace!("intra coded_block_pattern num: {} pattern: {:?}", x, result);
+    result
+}
+
+#[inline]
+pub fn code_num_to_inter_coded_block_pattern(x: u8) -> Option<CodedBlockPattern> {
+    let result =
+        CODE_NUM_TO_INTER_CODED_BLOCK_PATTERN.get(x as usize).map(|x| CodedBlockPattern(*x));
+    trace!("inter coded_block_pattern num: {} pattern: {:?}", x, result);
     result
 }