cros_codecs/codec/h264/
synthesizer.rs

1// Copyright 2024 The ChromiumOS Authors
2// Use of this source code is governed by a BSD-style license that can be
3// found in the LICENSE file.
4use std::io::Write;
5
6use thiserror::Error;
7
8use crate::codec::h264::nalu_writer::NaluWriter;
9use crate::codec::h264::nalu_writer::NaluWriterError;
10use crate::codec::h264::parser::HrdParams;
11use crate::codec::h264::parser::NaluType;
12use crate::codec::h264::parser::Pps;
13use crate::codec::h264::parser::Sps;
14use crate::codec::h264::parser::DEFAULT_4X4_INTER;
15use crate::codec::h264::parser::DEFAULT_4X4_INTRA;
16use crate::codec::h264::parser::DEFAULT_8X8_INTER;
17use crate::codec::h264::parser::DEFAULT_8X8_INTRA;
18
19mod private {
20    pub trait NaluStruct {}
21}
22
23impl private::NaluStruct for Sps {}
24
25impl private::NaluStruct for Pps {}
26
27#[derive(Error, Debug)]
28pub enum SynthesizerError {
29    #[error("tried to synthesize unsupported settings")]
30    Unsupported,
31    #[error(transparent)]
32    NaluWriter(#[from] NaluWriterError),
33}
34
35pub type SynthesizerResult<T> = Result<T, SynthesizerError>;
36
37/// A helper to output typed NALUs to [`std::io::Write`] using [`NaluWriter`].
38pub struct Synthesizer<'n, N: private::NaluStruct, W: Write> {
39    writer: NaluWriter<W>,
40    nalu: &'n N,
41}
42
43/// Extended Sample Aspect Ratio - H.264 Table E-1
44const EXTENDED_SAR: u8 = 255;
45
46impl<N: private::NaluStruct, W: Write> Synthesizer<'_, N, W> {
47    fn u<T: Into<u32>>(&mut self, bits: usize, value: T) -> SynthesizerResult<()> {
48        self.writer.write_u(bits, value)?;
49        Ok(())
50    }
51
52    fn f<T: Into<u32>>(&mut self, bits: usize, value: T) -> SynthesizerResult<()> {
53        self.writer.write_f(bits, value)?;
54        Ok(())
55    }
56
57    fn ue<T: Into<u32>>(&mut self, value: T) -> SynthesizerResult<()> {
58        self.writer.write_ue(value)?;
59        Ok(())
60    }
61
62    fn se<T: Into<i32>>(&mut self, value: T) -> SynthesizerResult<()> {
63        self.writer.write_se(value)?;
64        Ok(())
65    }
66
67    fn scaling_list(&mut self, list: &[u8], default: &[u8]) -> SynthesizerResult<()> {
68        // H.264 7.3.2.1.1.1
69        if list == default {
70            self.se(-8)?;
71            return Ok(());
72        }
73
74        // The number of list values we want to encode.
75        let mut run = list.len();
76
77        // Check how many values at the end of the matrix are the same,
78        // so we can save on encoding those.
79        for j in (1..list.len()).rev() {
80            if list[j - 1] != list[j] {
81                break;
82            }
83            run -= 1;
84        }
85
86        // Encode deltas.
87        let mut last_scale = 8;
88        for scale in &list[0..run] {
89            let delta_scale = *scale as i32 - last_scale;
90            self.se(delta_scale)?;
91            last_scale = *scale as i32;
92        }
93
94        // Didn't encode all values, encode -|last_scale| to set decoder's
95        // |next_scale| (H.264 7.3.2.1.1.1) to zero, i.e. decoder should repeat
96        // last values in matrix.
97        if run < list.len() {
98            self.se(-last_scale)?;
99        }
100
101        Ok(())
102    }
103
104    fn default_scaling_list(i: usize) -> &'static [u8] {
105        // H.264 Table 7-2
106        match i {
107            0 => &DEFAULT_4X4_INTRA[..],
108            1 => &DEFAULT_4X4_INTRA[..],
109            2 => &DEFAULT_4X4_INTRA[..],
110            3 => &DEFAULT_4X4_INTER[..],
111            4 => &DEFAULT_4X4_INTER[..],
112            5 => &DEFAULT_4X4_INTER[..],
113            6 => &DEFAULT_8X8_INTRA[..],
114            7 => &DEFAULT_8X8_INTER[..],
115            8 => &DEFAULT_8X8_INTRA[..],
116            9 => &DEFAULT_8X8_INTER[..],
117            10 => &DEFAULT_8X8_INTRA[..],
118            11 => &DEFAULT_8X8_INTER[..],
119            _ => unreachable!(),
120        }
121    }
122
123    fn rbsp_trailing_bits(&mut self) -> SynthesizerResult<()> {
124        self.f(1, 1u32)?;
125
126        while !self.writer.aligned() {
127            self.f(1, 0u32)?;
128        }
129
130        Ok(())
131    }
132}
133
134impl<'n, W: Write> Synthesizer<'n, Sps, W> {
135    pub fn synthesize(
136        ref_idc: u8,
137        sps: &'n Sps,
138        writer: W,
139        ep_enabled: bool,
140    ) -> SynthesizerResult<()> {
141        let mut s = Self {
142            writer: NaluWriter::<W>::new(writer, ep_enabled),
143            nalu: sps,
144        };
145
146        s.writer.write_header(ref_idc, NaluType::Sps as u8)?;
147        s.seq_parameter_set_data()?;
148        s.rbsp_trailing_bits()
149    }
150
151    fn hrd_parameters(&mut self, hrd_params: &HrdParams) -> SynthesizerResult<()> {
152        self.ue(hrd_params.cpb_cnt_minus1)?;
153        self.u(4, hrd_params.bit_rate_scale)?;
154        self.u(4, hrd_params.cpb_size_scale)?;
155
156        for i in 0..=(hrd_params.cpb_cnt_minus1 as usize) {
157            self.ue(hrd_params.bit_rate_value_minus1[i])?;
158            self.ue(hrd_params.cpb_size_value_minus1[i])?;
159            self.u(1, hrd_params.cbr_flag[i])?;
160        }
161
162        self.u(5, hrd_params.initial_cpb_removal_delay_length_minus1)?;
163        self.u(5, hrd_params.cpb_removal_delay_length_minus1)?;
164        self.u(5, hrd_params.dpb_output_delay_length_minus1)?;
165        self.u(5, hrd_params.time_offset_length)?;
166
167        Ok(())
168    }
169
170    fn vui_parameters(&mut self) -> SynthesizerResult<()> {
171        // H.264 E.1.1
172        let vui_params = &self.nalu.vui_parameters;
173
174        self.u(1, vui_params.aspect_ratio_info_present_flag)?;
175        if vui_params.aspect_ratio_info_present_flag {
176            self.u(8, vui_params.aspect_ratio_idc)?;
177            if vui_params.aspect_ratio_idc == EXTENDED_SAR {
178                self.u(16, vui_params.sar_width)?;
179                self.u(16, vui_params.sar_height)?;
180            }
181        }
182
183        self.u(1, vui_params.overscan_info_present_flag)?;
184        if vui_params.overscan_info_present_flag {
185            self.u(1, vui_params.overscan_appropriate_flag)?;
186        }
187
188        self.u(1, vui_params.video_signal_type_present_flag)?;
189        if vui_params.video_signal_type_present_flag {
190            self.u(3, vui_params.video_format)?;
191            self.u(1, vui_params.video_full_range_flag)?;
192
193            self.u(1, vui_params.colour_description_present_flag)?;
194            if vui_params.colour_description_present_flag {
195                self.u(8, vui_params.colour_primaries)?;
196                self.u(8, vui_params.transfer_characteristics)?;
197                self.u(8, vui_params.matrix_coefficients)?;
198            }
199        }
200
201        self.u(1, vui_params.chroma_loc_info_present_flag)?;
202        if vui_params.chroma_loc_info_present_flag {
203            self.ue(vui_params.chroma_sample_loc_type_top_field)?;
204            self.ue(self.nalu.vui_parameters.chroma_sample_loc_type_bottom_field)?;
205        }
206
207        self.u(1, vui_params.timing_info_present_flag)?;
208        if vui_params.timing_info_present_flag {
209            self.u(32, vui_params.num_units_in_tick)?;
210            self.u(32, vui_params.time_scale)?;
211            self.u(1, vui_params.fixed_frame_rate_flag)?;
212        }
213
214        self.u(1, vui_params.nal_hrd_parameters_present_flag)?;
215        if vui_params.nal_hrd_parameters_present_flag {
216            self.hrd_parameters(&vui_params.nal_hrd_parameters)?;
217        }
218        self.u(1, vui_params.vcl_hrd_parameters_present_flag)?;
219        if vui_params.vcl_hrd_parameters_present_flag {
220            self.hrd_parameters(&vui_params.vcl_hrd_parameters)?;
221        }
222
223        if vui_params.nal_hrd_parameters_present_flag || vui_params.vcl_hrd_parameters_present_flag
224        {
225            self.u(1, vui_params.low_delay_hrd_flag)?;
226        }
227
228        self.u(1, vui_params.pic_struct_present_flag)?;
229
230        self.u(1, vui_params.bitstream_restriction_flag)?;
231        if vui_params.bitstream_restriction_flag {
232            self.u(1, vui_params.motion_vectors_over_pic_boundaries_flag)?;
233            self.ue(vui_params.max_bytes_per_pic_denom)?;
234            self.ue(vui_params.max_bits_per_mb_denom)?;
235            self.ue(vui_params.log2_max_mv_length_horizontal)?;
236            self.ue(vui_params.log2_max_mv_length_vertical)?;
237            self.ue(vui_params.max_num_reorder_frames)?;
238            self.ue(vui_params.max_dec_frame_buffering)?;
239        }
240
241        Ok(())
242    }
243
244    fn seq_parameter_set_data(&mut self) -> SynthesizerResult<()> {
245        // H.264 7.3.2.1.1
246        self.u(8, self.nalu.profile_idc)?;
247        self.u(1, self.nalu.constraint_set0_flag)?;
248        self.u(1, self.nalu.constraint_set1_flag)?;
249        self.u(1, self.nalu.constraint_set2_flag)?;
250        self.u(1, self.nalu.constraint_set3_flag)?;
251        self.u(1, self.nalu.constraint_set4_flag)?;
252        self.u(1, self.nalu.constraint_set5_flag)?;
253        self.u(2, /* reserved_zero_2bits */ 0u32)?;
254        self.u(8, self.nalu.level_idc as u32)?;
255        self.ue(self.nalu.seq_parameter_set_id)?;
256
257        if self.nalu.profile_idc == 100
258            || self.nalu.profile_idc == 110
259            || self.nalu.profile_idc == 122
260            || self.nalu.profile_idc == 244
261            || self.nalu.profile_idc == 44
262            || self.nalu.profile_idc == 83
263            || self.nalu.profile_idc == 86
264            || self.nalu.profile_idc == 118
265            || self.nalu.profile_idc == 128
266            || self.nalu.profile_idc == 138
267            || self.nalu.profile_idc == 139
268            || self.nalu.profile_idc == 134
269            || self.nalu.profile_idc == 135
270        {
271            self.ue(self.nalu.chroma_format_idc)?;
272
273            if self.nalu.chroma_format_idc == 3 {
274                self.u(1, self.nalu.separate_colour_plane_flag)?;
275            }
276
277            self.ue(self.nalu.bit_depth_luma_minus8)?;
278            self.ue(self.nalu.bit_depth_chroma_minus8)?;
279            self.u(1, self.nalu.qpprime_y_zero_transform_bypass_flag)?;
280            self.u(1, self.nalu.seq_scaling_matrix_present_flag)?;
281
282            if self.nalu.seq_scaling_matrix_present_flag {
283                let scaling_list_count = if self.nalu.chroma_format_idc != 3 {
284                    8
285                } else {
286                    12
287                };
288
289                for i in 0..scaling_list_count {
290                    // Assume if scaling lists are zeroed that they are not present.
291                    if i < 6 {
292                        if self.nalu.scaling_lists_4x4[i] == [0; 16] {
293                            self.u(1, /* seq_scaling_list_present_flag */ false)?;
294                        } else {
295                            self.u(1, /* seq_scaling_list_present_flag */ true)?;
296                            self.scaling_list(
297                                &self.nalu.scaling_lists_4x4[i],
298                                Self::default_scaling_list(i),
299                            )?;
300                        }
301                    } else if self.nalu.scaling_lists_8x8[i - 6] == [0; 64] {
302                        self.u(1, /* seq_scaling_list_present_flag */ false)?;
303                    } else {
304                        self.u(1, /* seq_scaling_list_present_flag */ true)?;
305                        self.scaling_list(
306                            &self.nalu.scaling_lists_8x8[i - 6],
307                            Self::default_scaling_list(i),
308                        )?;
309                    }
310                }
311            }
312        }
313
314        self.ue(self.nalu.log2_max_frame_num_minus4)?;
315        self.ue(self.nalu.pic_order_cnt_type)?;
316
317        if self.nalu.pic_order_cnt_type == 0 {
318            self.ue(self.nalu.log2_max_pic_order_cnt_lsb_minus4)?;
319        } else if self.nalu.pic_order_cnt_type == 1 {
320            self.u(1, self.nalu.delta_pic_order_always_zero_flag)?;
321            self.se(self.nalu.offset_for_non_ref_pic)?;
322            self.se(self.nalu.offset_for_top_to_bottom_field)?;
323            self.ue(self.nalu.num_ref_frames_in_pic_order_cnt_cycle)?;
324
325            for offset_for_ref_frame in &self.nalu.offset_for_ref_frame {
326                self.se(*offset_for_ref_frame)?;
327            }
328        }
329
330        self.ue(self.nalu.max_num_ref_frames)?;
331        self.u(1, self.nalu.gaps_in_frame_num_value_allowed_flag)?;
332        self.ue(self.nalu.pic_width_in_mbs_minus1)?;
333        self.ue(self.nalu.pic_height_in_map_units_minus1)?;
334        self.u(1, self.nalu.frame_mbs_only_flag)?;
335        if !self.nalu.frame_mbs_only_flag {
336            self.u(1, self.nalu.mb_adaptive_frame_field_flag)?;
337        }
338        self.u(1, self.nalu.direct_8x8_inference_flag)?;
339
340        self.u(1, self.nalu.frame_cropping_flag)?;
341        if self.nalu.frame_cropping_flag {
342            self.ue(self.nalu.frame_crop_left_offset)?;
343            self.ue(self.nalu.frame_crop_right_offset)?;
344            self.ue(self.nalu.frame_crop_top_offset)?;
345            self.ue(self.nalu.frame_crop_bottom_offset)?;
346        }
347
348        self.u(1, self.nalu.vui_parameters_present_flag)?;
349        if self.nalu.vui_parameters_present_flag {
350            self.vui_parameters()?;
351        }
352
353        Ok(())
354    }
355}
356
357impl<'n, W: Write> Synthesizer<'n, Pps, W> {
358    pub fn synthesize(
359        ref_idc: u8,
360        pps: &'n Pps,
361        writer: W,
362        ep_enabled: bool,
363    ) -> SynthesizerResult<()> {
364        let mut s = Self {
365            writer: NaluWriter::<W>::new(writer, ep_enabled),
366            nalu: pps,
367        };
368
369        s.writer.write_header(ref_idc, NaluType::Pps as u8)?;
370        s.pic_parameter_set_rbsp()?;
371        s.rbsp_trailing_bits()
372    }
373
374    fn pic_parameter_set_rbsp(&mut self) -> SynthesizerResult<()> {
375        self.ue(self.nalu.pic_parameter_set_id)?;
376        self.ue(self.nalu.seq_parameter_set_id)?;
377        self.u(1, self.nalu.entropy_coding_mode_flag)?;
378        self.u(1, self.nalu.bottom_field_pic_order_in_frame_present_flag)?;
379
380        self.ue(self.nalu.num_slice_groups_minus1)?;
381        if self.nalu.num_slice_groups_minus1 > 0 {
382            return Err(SynthesizerError::Unsupported);
383        }
384
385        self.ue(self.nalu.num_ref_idx_l0_default_active_minus1)?;
386        self.ue(self.nalu.num_ref_idx_l1_default_active_minus1)?;
387        self.u(1, self.nalu.weighted_pred_flag)?;
388        self.u(2, self.nalu.weighted_bipred_idc)?;
389        self.se(self.nalu.pic_init_qp_minus26)?;
390        self.se(self.nalu.pic_init_qs_minus26)?;
391        self.se(self.nalu.chroma_qp_index_offset)?;
392        self.u(1, self.nalu.deblocking_filter_control_present_flag)?;
393        self.u(1, self.nalu.constrained_intra_pred_flag)?;
394        self.u(1, self.nalu.redundant_pic_cnt_present_flag)?;
395
396        if !(self.nalu.transform_8x8_mode_flag
397            || self.nalu.pic_scaling_matrix_present_flag
398            || self.nalu.second_chroma_qp_index_offset != 0)
399        {
400            return Ok(());
401        }
402
403        self.u(1, self.nalu.transform_8x8_mode_flag)?;
404        self.u(1, self.nalu.pic_scaling_matrix_present_flag)?;
405
406        if self.nalu.pic_scaling_matrix_present_flag {
407            let mut scaling_list_count = 6;
408            if self.nalu.transform_8x8_mode_flag {
409                if self.nalu.sps.chroma_format_idc != 3 {
410                    scaling_list_count += 2;
411                } else {
412                    scaling_list_count += 6;
413                }
414            }
415
416            for i in 0..scaling_list_count {
417                // Assume if scaling lists are zeroed that they are not present.
418                if i < 6 {
419                    if self.nalu.scaling_lists_4x4[i] == [0; 16] {
420                        self.u(1, /* seq_scaling_list_present_flag */ false)?;
421                    } else {
422                        self.u(1, /* seq_scaling_list_present_flag */ true)?;
423                        self.scaling_list(
424                            &self.nalu.scaling_lists_4x4[i],
425                            Self::default_scaling_list(i),
426                        )?;
427                    }
428                } else if self.nalu.scaling_lists_8x8[i - 6] == [0; 64] {
429                    self.u(1, /* seq_scaling_list_present_flag */ false)?;
430                } else {
431                    self.u(1, /* seq_scaling_list_present_flag */ true)?;
432                    self.scaling_list(
433                        &self.nalu.scaling_lists_8x8[i - 6],
434                        Self::default_scaling_list(i),
435                    )?;
436                }
437            }
438        }
439
440        self.se(self.nalu.second_chroma_qp_index_offset)?;
441
442        Ok(())
443    }
444}
445
446#[cfg(test)]
447mod tests {
448    use std::io::Cursor;
449
450    use super::*;
451    use crate::codec::h264::parser::Nalu;
452    use crate::codec::h264::parser::NaluType;
453    use crate::codec::h264::parser::Parser;
454    use crate::codec::h264::parser::Profile;
455
456    #[test]
457    fn synthesize_sps() {
458        let raw_sps_buf = [0x00, 0x00, 0x00, 0x01, 0x07, 0x00, 0x00, 0x0a, 0xfb, 0x88];
459        let mut raw_sps = Cursor::new(&raw_sps_buf[..]);
460
461        let nalu = Nalu::next(&mut raw_sps).unwrap();
462        assert_eq!(nalu.header.type_, NaluType::Sps);
463
464        let mut parser = Parser::default();
465        let sps = parser.parse_sps(&nalu).unwrap();
466
467        let mut buf = Vec::<u8>::new();
468        Synthesizer::<'_, Sps, _>::synthesize(0, sps, &mut buf, false).unwrap();
469
470        assert_eq!(buf, raw_sps_buf);
471
472        let write_to_file = std::option_env!("CROS_CODECS_TEST_WRITE_TO_FILE") == Some("true");
473        if write_to_file {
474            let mut out = std::fs::File::create("sps.h264").unwrap();
475            out.write_all(&buf).unwrap();
476            out.flush().unwrap();
477        }
478
479        let mut cursor = Cursor::new(&buf[..]);
480        let nalu = Nalu::next(&mut cursor).unwrap();
481
482        let mut parser = Parser::default();
483
484        let sps2 = parser.parse_sps(&nalu).unwrap();
485
486        assert_eq!(sps, sps2);
487    }
488
489    #[test]
490    fn synthesize_sps_scaling_lists() {
491        let sps = Sps {
492            profile_idc: Profile::High as u8,
493            seq_scaling_matrix_present_flag: true,
494            scaling_lists_4x4: [[
495                11, 20, 10, 20, 10, 22, 10, 20, 10, 20, 13, 20, 10, 20, 10, 24,
496            ]; 6],
497            scaling_lists_8x8: [
498                [
499                    33, 20, 10, 21, 33, 20, 12, 20, 33, 23, 10, 20, 33, 20, 10, 20, 33, 24, 10, 20,
500                    33, 20, 15, 20, 33, 20, 10, 26, 33, 20, 17, 20, 33, 28, 10, 20, 33, 20, 10, 20,
501                    33, 29, 10, 20, 33, 20, 11, 20, 33, 20, 10, 20, 33, 20, 10, 20, 33, 20, 10, 20,
502                    33, 20, 10, 20,
503                ],
504                [
505                    10, 77, 11, 20, 10, 77, 12, 20, 10, 77, 13, 20, 10, 77, 14, 20, 10, 77, 15, 20,
506                    10, 77, 16, 20, 10, 77, 17, 20, 10, 77, 18, 20, 10, 77, 19, 20, 10, 77, 10, 20,
507                    10, 77, 10, 21, 10, 77, 10, 22, 10, 77, 10, 23, 10, 77, 10, 24, 10, 77, 10, 26,
508                    10, 77, 10, 28,
509                ],
510                [0; 64],
511                [0; 64],
512                [0; 64],
513                [0; 64],
514            ],
515            frame_mbs_only_flag: true,
516            ..Default::default()
517        };
518
519        let mut buf = Vec::<u8>::new();
520        Synthesizer::<'_, Sps, _>::synthesize(0, &sps, &mut buf, false).unwrap();
521
522        let write_to_file = std::option_env!("CROS_CODECS_TEST_WRITE_TO_FILE") == Some("true");
523        if write_to_file {
524            let mut out = std::fs::File::create("sps.h264").unwrap();
525            out.write_all(&buf).unwrap();
526            out.flush().unwrap();
527        }
528
529        let mut cursor = Cursor::new(&buf[..]);
530        let nalu = Nalu::next(&mut cursor).unwrap();
531
532        let mut parser = Parser::default();
533
534        let sps2 = parser.parse_sps(&nalu).unwrap();
535
536        assert_eq!(sps.scaling_lists_4x4, sps2.scaling_lists_4x4);
537        assert_eq!(sps.scaling_lists_8x8, sps2.scaling_lists_8x8);
538    }
539
540    #[test]
541    fn synthesize_pps() {
542        let raw_sps_pps = [
543            0x00, 0x00, 0x00, 0x01, 0x07, 0x4d, 0x40, 0x0d, 0xa9, 0x18, 0x28, 0x3e, 0x60, 0x0d,
544            0x41, 0x80, 0x41, 0xad, 0xb0, 0xad, 0x7b, 0xdf, 0x01, 0x00, 0x00, 0x00, 0x01, 0x08,
545            0xde, 0x09, 0x88,
546        ];
547
548        let mut buf = Vec::<u8>::new();
549        let mut out = Cursor::new(&mut buf);
550
551        let mut cursor = Cursor::new(&raw_sps_pps[..]);
552        let mut parser: Parser = Default::default();
553
554        while let Ok(nalu) = Nalu::next(&mut cursor) {
555            match nalu.header.type_ {
556                NaluType::Sps => {
557                    let sps = parser.parse_sps(&nalu).unwrap();
558                    Synthesizer::<'_, Sps, _>::synthesize(0, sps, &mut out, false).unwrap();
559                }
560                NaluType::Pps => {
561                    let pps = parser.parse_pps(&nalu).unwrap();
562                    Synthesizer::<'_, Pps, _>::synthesize(0, pps, &mut out, false).unwrap();
563                }
564                _ => panic!(),
565            }
566        }
567
568        let write_to_file = std::option_env!("CROS_CODECS_TEST_WRITE_TO_FILE") == Some("true");
569        if write_to_file {
570            let mut out = std::fs::File::create("sps_pps.h264").unwrap();
571            out.write_all(&buf).unwrap();
572            out.flush().unwrap();
573
574            let mut out = std::fs::File::create("sps_pps_ref.h264").unwrap();
575            out.write_all(&raw_sps_pps).unwrap();
576            out.flush().unwrap();
577        }
578
579        assert_eq!(buf, raw_sps_pps);
580    }
581}
cros_codecs/codec/h264/synthesizer.rs

cros_codecs/codec/h264/
synthesizer.rs
