decode and compare all frames in the video, not just the first one

Author: Djuffin

src/h264/decoder.rs

@@ -59,7 +59,7 @@ impl DecoderContext {
 #[derive(Debug, Default)]
 pub struct Decoder {
     context: DecoderContext,
-    frame_buffer: Option<VideoFrame>,
+    frame_buffer: Vec<VideoFrame>,
 }
 
 impl Decoder {
@@ -108,8 +108,17 @@ impl Decoder {
 
                     info!("non-IDR Slice: {:#?}", slice);
                     if slice.header.slice_type != SliceType::I {
-                        break;
+                        todo!("implement P-slices")
                     }
+
+                    let frame = VideoFrame::new_with_padding(
+                        slice.sps.pic_width(),
+                        slice.sps.pic_hight(),
+                        v_frame::pixel::ChromaSampling::Cs420,
+                        16,
+                    );
+                    self.frame_buffer.push(frame);
+
                     parser::parse_slice_data(&mut unit_input, &mut slice)
                         .map_err(parse_error_handler)?;
                     info!("Blocks: {:#?}", slice.get_macroblock_count());
@@ -121,23 +130,24 @@ impl Decoder {
                             .map_err(parse_error_handler)?;
 
                     info!("IDR Slice: {:#?}", slice);
+                    let frame = VideoFrame::new_with_padding(
+                        slice.sps.pic_width(),
+                        slice.sps.pic_hight(),
+                        v_frame::pixel::ChromaSampling::Cs420,
+                        16,
+                    );
+                    self.frame_buffer.push(frame);
+
                     parser::parse_slice_data(&mut unit_input, &mut slice)
                         .map_err(parse_error_handler)?;
                     info!("Blocks: {:#?}", slice.get_macroblock_count());
-                    return self.process_slice(&mut slice); // Temporarily stop after first slice
+                    self.process_slice(&mut slice)?; // Temporarily stop after first slice
                 }
                 NalUnitType::SupplementalEnhancementInfo => {}
                 NalUnitType::SeqParameterSet => {
                     let sps = parser::parse_sps(&mut unit_input).map_err(parse_error_handler)?;
                     info!("SPS: {:#?}", sps);
                     assert_eq!(sps.ChromaArrayType(), ChromaFormat::YUV420);
-                    let frame = VideoFrame::new_with_padding(
-                        sps.pic_width(),
-                        sps.pic_hight(),
-                        v_frame::pixel::ChromaSampling::Cs420,
-                        16,
-                    );
-                    self.frame_buffer = Some(frame);
                     self.context.put_sps(sps);
                 }
                 NalUnitType::PicParameterSet => {
@@ -162,16 +172,16 @@ impl Decoder {
         Ok(())
     }
 
-    pub fn get_frame_buffer(&self) -> Option<&VideoFrame> {
-        self.frame_buffer.as_ref()
+    pub fn get_frame_buffer(&self) -> &[VideoFrame] {
+        self.frame_buffer.as_slice()
     }
 
     fn process_slice(&mut self, slice: &mut Slice) -> Result<(), DecodingError> {
-        if self.frame_buffer.is_none() {
+        if self.frame_buffer.is_empty() {
             return Err(DecodingError::Wtf);
         }
         let mut qp = slice.pps.pic_init_qp_minus26 + 26 + slice.header.slice_qp_delta;
-        let frame = self.frame_buffer.as_mut().unwrap();
+        let frame = self.frame_buffer.last_mut().unwrap();
         for mb_addr in 0..(slice.sps.pic_size_in_mbs() as u32) {
             let mb_loc = slice.get_mb_location(mb_addr);
             if let Some(mb) = slice.get_mb(mb_addr) {

src/main.rs

@@ -17,7 +17,7 @@ use std::io;
 
 use log::info;
 use v_frame::plane::PlaneOffset;
-use y4m_cmp::compare_frames;
+use y4m_cmp::{compare_files, compare_frames};
 
 fn main() {
     diag::init(false);
@@ -30,39 +30,38 @@ fn main() {
     let buf = fs::read(input_filename).unwrap();
     let mut decoder = h264::decoder::Decoder::new();
     decoder.decode(&buf).expect("parsing error");
-    let frame = decoder.get_frame_buffer().unwrap();
-    let y_plane = &frame.planes[0];
 
+    let first_frame = decoder.get_frame_buffer().first().unwrap();
+    let y_plane = &first_frame.planes[0];
     let w = y_plane.cfg.width as u32;
     let h = y_plane.cfg.height as u32;
 
-    info!("Writing frame {w} x {h} to y4m");
     let mut writer = io::BufWriter::new(fs::File::create("output.y4m").unwrap());
     let mut encoder = y4m::encode(w as usize, h as usize, y4m::Ratio { num: 15, den: 1 })
         .with_colorspace(y4m::Colorspace::C420)
         .write_header(&mut writer)
         .unwrap();
-    let mut planes = Vec::<Vec<u8>>::new();
-    for plane in &frame.planes {
-        let data_size = plane.cfg.width * plane.cfg.height;
-        let mut data = vec![0; data_size];
-        plane.copy_to_raw_u8(&mut data, plane.cfg.width, 1);
-        planes.push(data)
-    }
 
-    let yuv_frame =
-        y4m::Frame::new([planes[0].as_slice(), planes[1].as_slice(), planes[2].as_slice()], None);
+    for (num, frame) in decoder.get_frame_buffer().iter().enumerate() {
+        info!("Writing frame #{num} {w} x {h} to y4m");
+
+        let mut planes = Vec::<Vec<u8>>::new();
+        for plane in &frame.planes {
+            let data_size = plane.cfg.width * plane.cfg.height;
+            let mut data = vec![0; data_size];
+            plane.copy_to_raw_u8(&mut data, plane.cfg.width, 1);
+            planes.push(data)
+        }
+
+        let yuv_frame = y4m::Frame::new(
+            [planes[0].as_slice(), planes[1].as_slice(), planes[2].as_slice()],
+            None,
+        );
+        encoder.write_frame(&yuv_frame).unwrap();
+    }
+    drop(encoder);
 
-    if let Ok(expected_file) = fs::File::open("data/NL1_Sony_D.y4m") {
-        let reader = io::BufReader::new(expected_file);
-        let mut decoder = y4m::Decoder::new(reader).unwrap();
-        let expected_h = decoder.get_height();
-        let expected_w = decoder.get_width();
-        let expected_frame = decoder.read_frame().unwrap();
-        assert!(w == expected_w as u32);
-        assert!(h == expected_h as u32);
-        let compare_result = compare_frames(expected_w, expected_h, &yuv_frame, &expected_frame);
-        print!("Frame comparison result: {compare_result}");
+    if let Err(message) = compare_files("output.y4m", "data/NL1_Sony_D.y4m") {
+        print!("File comparison result: {message}");
     }
-    encoder.write_frame(&yuv_frame).unwrap();
 }

src/y4m_cmp.rs

@@ -1,5 +1,12 @@
+use v_frame::frame;
+
 use crate::h264::tables::MB_WIDTH;
 
+use std::fmt::format;
+use std::fs;
+use std::io;
+use std::io::Error;
+
 fn compare_plane(
     width: usize,
     height: usize,
@@ -54,3 +61,46 @@ pub fn compare_frames(
 
     result
 }
+
+pub fn compare_files(actual_filename: &str, expected_filename: &str) -> Result<(), String> {
+    fn stringify(x: Error) -> String {
+        format!("error code: {x}")
+    }
+    fn y4m_stringify(x: y4m::Error) -> String {
+        format!("error code: {x}")
+    }
+
+    let expected_file = fs::File::open(expected_filename).map_err(stringify)?;
+    let expected_reader = io::BufReader::new(expected_file);
+    let mut expected_decoder = y4m::Decoder::new(expected_reader).map_err(y4m_stringify)?;
+
+    let actual_file = fs::File::open(actual_filename).map_err(stringify)?;
+    let actual_reader = io::BufReader::new(actual_file);
+    let mut actual_decoder = y4m::Decoder::new(actual_reader).map_err(y4m_stringify)?;
+
+    let expected_h = expected_decoder.get_height();
+    let expected_w = expected_decoder.get_width();
+    let actual_h = actual_decoder.get_height();
+    let actual_w = actual_decoder.get_width();
+    if (expected_w, expected_h) != (actual_w, actual_h) {
+        return Err(format!("Unexpected size of frames. {actual_w}x{actual_h} vs expected {expected_w}x{expected_h}"));
+    }
+
+    let mut frame_idx = 0;
+    while let (Ok(actual_frame), Ok(expected_frame)) =
+        (actual_decoder.read_frame(), expected_decoder.read_frame())
+    {
+        let compare_result = compare_frames(expected_w, expected_h, &actual_frame, &expected_frame);
+        if !compare_result.is_empty() {
+            return Err(format!("Frame #{frame_idx} mismatch: {compare_result}"));
+        }
+        frame_idx += 1;
+    }
+
+    if let Err(y4m::Error::EOF) = actual_decoder.read_frame() {
+    } else {
+        return Err(format!("Unexpected number of frames. {frame_idx}"));
+    }
+
+    Ok(())
+}