import logging from dataclasses import dataclass, field from math import floor, log, ceil from typeguard import typechecked import hexdump from tscut.h264.bitstream import ( read_bit, read_bits, read_boolean, read_byte, read_long, read_signed_exp_golomb, read_unsigned_exp_golomb, read_word, write_bit, write_bits, write_boolean, write_byte, write_long, write_rbsp_trailing_bits, write_scaling_list, write_signed_exp_golomb, write_unsigned_exp_golomb, write_word, parse_scaling_list, parse_rbsp_trailing_bits, more_rbsp_data ) logger = logging.getLogger(__name__) @dataclass class HRD: cpb_cnt_minus1: int=0 bit_rate_scale: int=0 cpb_size_scale: int=0 bit_rate_value_minus1: dict = field(default_factory=dict) cpb_size_value_minus1: dict = field(default_factory=dict) cbr_flag: dict = field(default_factory=dict) initial_cpb_removal_delay_length_minus1: int=0 cpb_removal_delay_length_minus1: int=0 dpb_output_delay_length_minus1: int=0 time_offset_length: int=0 def __init__(self): self.bit_rate_value_minus1 = {} self.cpb_size_value_minus1 = {} self.cbr_flag = {} def from_bytes(self, buf, bit_position): bit_position, self.cpb_cnt_minus1 = read_unsigned_exp_golomb(buf, bit_position) bit_position, self.bit_rate_scale = read_bits(buf, bit_position, 4) bit_position, self.cpb_size_scale = read_bits(buf, bit_position, 4) for i in range(self.cpb_cnt_minus1+1): bit_position, v = read_unsigned_exp_golomb(buf, bit_position) self.bit_rate_value_minus1[i] = v bit_position, v = read_unsigned_exp_golomb(buf, bit_position) self.cpb_size_value_minus1[i] = v bit_position, b = read_boolean(buf, bit_position) self.cbr_flag[i] = b bit_position, self.initial_cpb_removal_delay_length_minus1 = read_bits(buf, bit_position, 5) bit_position, self.cpb_removal_delay_length_minus1 = read_bits(buf, bit_position, 5) bit_position, self.dpb_output_delay_length_minus1 = read_bits(buf, bit_position, 5) bit_position, self.time_offset_length = read_bits(buf, bit_position, 5) return bit_position def to_bytes(self, buf, bit_position): bit_position = write_unsigned_exp_golomb(buf, bit_position, self.cpb_cnt_minus1) bit_position = write_bits(buf, bit_position, self.bit_rate_scale, 4) bit_position = write_bits(buf, bit_position, self.cpb_size_scale, 4) for i in range(self.cpb_cnt_minus1+1): v = self.bit_rate_value_minus1[i] bit_position = write_unsigned_exp_golomb(buf, bit_position, v) v = self.cpb_size_value_minus1[i] bit_position = write_unsigned_exp_golomb(buf, bit_position, v) b = self.cbr_flag[i] bit_position = write_boolean(buf, bit_position, b) bit_position = write_bits(buf, bit_position, self.initial_cpb_removal_delay_length_minus1, 5) bit_position = write_bits(buf, bit_position, self.cpb_removal_delay_length_minus1, 5) bit_position = write_bits(buf, bit_position, self.dpb_output_delay_length_minus1, 5) bit_position = write_bits(buf, bit_position, self.time_offset_length, 5) return bit_position @dataclass class VUI: aspect_ratio_info_present_flag:bool=False aspect_ratio_idc:int=0 sar_width:int=0 sar_height:int=0 overscan_info_present_flag:bool=False overscan_appropriate_flag:bool=False video_signal_type_present_flag:bool=False video_format:int=0 video_full_range_flag:bool=False colour_description_present_flag:bool=False colour_primaries:int=0 transfer_characteristics:int=0 matrix_coefficients:int=0 chroma_loc_info_present_flag:bool=False chroma_sample_loc_type_top_field:int=0 chroma_sample_loc_type_bottom_field:int=0 timing_info_present_flag:bool=False num_units_in_tick:int=0 time_scale:int=0 fixed_frame_rate_flag:bool=False nal_hrd_parameters_present_flag:bool=False hrd_parameters:HRD|None=None vcl_hrd_parameters_present_flag:bool=False vcl_hrd_parameters:HRD|None=None low_delay_hrd_flag:bool=False pic_struct_present_flag:bool=False bitstream_restriction_flag:bool=False motion_vectors_over_pic_boundaries_flag:bool=False max_bytes_per_pic_denom:int=0 max_bits_per_mb_denom:int=0 log2_max_mv_length_horizontal:int=0 log2_max_mv_length_vertical:int=0 max_num_reorder_frames:int=0 max_dec_frame_buffering:int=0 # This structure is not guaranteed to be located at a byte boundary. # We must explicitely indicate bit offset. def from_bytes(self, buf, bit_position): bit_position, self.aspect_ratio_info_present_flag = read_boolean(buf, bit_position) if self.aspect_ratio_info_present_flag: bit_position, self.aspect_ratio_idc = read_byte(buf, bit_position) if self.aspect_ratio_idc == 255: # Extended_SAR bit_position, self.sar_width = read_word(buf, bit_position) bit_position, self.sar_height = read_word(buf, bit_position) bit_position, self.overscan_info_present_flag = read_boolean(buf, bit_position) if self.overscan_info_present_flag: bit_position, self.overscan_appropriate_flag = read_boolean(buf, bit_position) bit_position, self.video_signal_type_present_flag = read_boolean(buf, bit_position) if self.video_signal_type_present_flag: bit_position, self.video_format = read_bits(buf, bit_position, 3) bit_position, self.video_full_range_flag = read_boolean(buf, bit_position) bit_position, self.colour_description_present_flag = read_boolean(buf, bit_position) if self.colour_description_present_flag: bit_position, self.colour_primaries = read_byte(buf, bit_position) bit_position, self.transfer_characteristics = read_byte(buf, bit_position) bit_position, self.matrix_coefficients = read_byte(buf, bit_position) bit_position, self.chroma_loc_info_present_flag = read_boolean(buf, bit_position) if self.chroma_loc_info_present_flag: bit_position, self.chroma_sample_loc_type_top_field =\ read_unsigned_exp_golomb(buf, bit_position) bit_position, self.chroma_sample_loc_type_bottom_field =\ read_unsigned_exp_golomb(buf,bit_position) bit_position, self.timing_info_present_flag = read_boolean(buf, bit_position) if self.timing_info_present_flag: bit_position, self.num_units_in_tick = read_long(buf, bit_position) bit_position, self.time_scale = read_long(buf, bit_position) bit_position, self.fixed_frame_rate_flag = read_boolean(buf, bit_position) bit_position, self.nal_hrd_parameters_present_flag = read_boolean(buf, bit_position) if self.nal_hrd_parameters_present_flag: hrd = HRD() bit_position = hrd.from_bytes(buf, bit_position) self.hrd_parameters = hrd bit_position, self.vcl_hrd_parameters_present_flag = read_boolean(buf, bit_position) if self.vcl_hrd_parameters_present_flag: hrd = HRD() bit_position = hrd.from_bytes(buf, bit_position) self.vcl_hrd_parameters = hrd if self.nal_hrd_parameters_present_flag or self.vcl_hrd_parameters_present_flag: bit_position, self.low_delay_hrd_flag = read_boolean(buf, bit_position) bit_position, self.pic_struct_present_flag = read_boolean(buf, bit_position) bit_position, self.bitstream_restriction_flag = read_boolean(buf, bit_position) if self.bitstream_restriction_flag: bit_position, self.motion_vectors_over_pic_boundaries_flag =\ read_boolean(buf, bit_position) bit_position, self.max_bytes_per_pic_denom = read_unsigned_exp_golomb(buf, bit_position) bit_position, self.max_bits_per_mb_denom = read_unsigned_exp_golomb(buf, bit_position) bit_position, self.log2_max_mv_length_horizontal = read_unsigned_exp_golomb(buf, bit_position) bit_position, self.log2_max_mv_length_vertical = read_unsigned_exp_golomb(buf, bit_position) bit_position, self.max_num_reorder_frames = read_unsigned_exp_golomb(buf, bit_position) bit_position, self.max_dec_frame_buffering = read_unsigned_exp_golomb(buf, bit_position) return bit_position def to_bytes(self, buf, bit_position): bit_position = write_boolean(buf, bit_position, self.aspect_ratio_info_present_flag) if self.aspect_ratio_info_present_flag: bit_position = write_byte(buf, bit_position, self.aspect_ratio_idc) if self.aspect_ratio_idc == 255: # Extended_SAR bit_position = write_word(buf, bit_position, self.sar_width) bit_position = write_word(buf, bit_position, self.sar_height) bit_position = write_boolean(buf, bit_position, self.overscan_info_present_flag) if self.overscan_info_present_flag: bit_position = write_boolean(buf, bit_position, self.overscan_appropriate_flag) bit_position = write_boolean(buf, bit_position, self.video_signal_type_present_flag) if self.video_signal_type_present_flag: bit_position = write_bits(buf, bit_position, self.video_format, 3) bit_position = write_boolean(buf, bit_position, self.video_full_range_flag) bit_position = write_boolean(buf, bit_position, self.colour_description_present_flag) if self.colour_description_present_flag: bit_position = write_byte(buf, bit_position, self.colour_primaries) bit_position = write_byte(buf, bit_position, self.transfer_characteristics) bit_position = write_byte(buf, bit_position, self.matrix_coefficients) bit_position = write_boolean(buf, bit_position, self.chroma_loc_info_present_flag) if self.chroma_loc_info_present_flag: bit_position = write_unsigned_exp_golomb(buf, bit_position, self.chroma_sample_loc_type_top_field) bit_position = write_unsigned_exp_golomb(buf, bit_position, self.chroma_sample_loc_type_bottom_field) bit_position = write_boolean(buf, bit_position, self.timing_info_present_flag ) if self.timing_info_present_flag: bit_position = write_long(buf, bit_position, self.num_units_in_tick ) bit_position = write_long(buf, bit_position, self.time_scale) bit_position = write_boolean(buf, bit_position, self.fixed_frame_rate_flag) bit_position = write_boolean(buf, bit_position, self.nal_hrd_parameters_present_flag) if self.nal_hrd_parameters_present_flag: bit_position = self.hrd_parameters.to_bytes(buf, bit_position) bit_position = write_boolean(buf, bit_position, self.vcl_hrd_parameters_present_flag) if self.vcl_hrd_parameters_present_flag: bit_position = self.vcl_hrd_parameters.to_bytes(buf, bit_position) if self.nal_hrd_parameters_present_flag or self.vcl_hrd_parameters_present_flag: bit_position = write_boolean(buf, bit_position, self.low_delay_hrd_flag) bit_position = write_boolean(buf, bit_position, self.pic_struct_present_flag) bit_position = write_boolean(buf, bit_position, self.bitstream_restriction_flag) if self.bitstream_restriction_flag: bit_position = write_boolean(buf, bit_position, self.motion_vectors_over_pic_boundaries_flag) bit_position = write_unsigned_exp_golomb(buf, bit_position, self.max_bytes_per_pic_denom) bit_position = write_unsigned_exp_golomb(buf, bit_position, self.max_bits_per_mb_denom) bit_position = write_unsigned_exp_golomb(buf, bit_position, self.log2_max_mv_length_horizontal) bit_position = write_unsigned_exp_golomb(buf, bit_position, self.log2_max_mv_length_vertical) bit_position = write_unsigned_exp_golomb(buf, bit_position, self.max_num_reorder_frames) bit_position = write_unsigned_exp_golomb(buf, bit_position, self.max_dec_frame_buffering) return bit_position @dataclass class SPS: profile_idc:int=0 # u(8) constraint_set0_flag:bool=False # u(1) constraint_set1_flag:bool=False # u(1) constraint_set2_flag:bool=False # u(1) constraint_set3_flag:bool=False # u(1) constraint_set4_flag:bool=False # u(1) constraint_set5_flag:bool=False # u(1) level_idc:int=0 # u(8) seq_parameter_set_id:int=0 # ue(v) chroma_format_idc:int=0 # ue(v) separate_colour_plane_flag:bool=False # u(1) bit_depth_luma_minus8:int=0 # ue(v) bit_depth_chroma_minus8:int=0 # ue(v) qpprime_y_zero_transform_bypass_flag:bool=False # u(1) seq_scaling_matrix_present_flag:bool=False # u(1) scaling_list: dict = field(default_factory=dict) log2_max_frame_num_minus4:int=0 # ue(v) pic_order_cnt_type:int=0 # ue(v) log2_max_pic_order_cnt_lsb_minus4:int=0 # ue(v) delta_pic_order_always_zero_flag:bool=False # ue(1) offset_for_non_ref_pic:int=0 # se(v) offset_for_top_to_bottom_field:int=0 # se(v) num_ref_frames_in_pic_order_cnt_cycle:int=0 # ue(v) offset_for_ref_frame:dict[int,int] = field(default_factory=dict) max_num_ref_frames:int=9 # ue(v) gaps_in_frame_num_value_allowed_flag:bool=False # u(1) pic_width_in_mbs_minus1:int=0 # ue(v) pic_height_in_map_units_minus1:int=0 # ue(v) frame_mbs_only_flag:bool=False # u(1) mb_adaptive_frame_field_flag:bool=False # u(1) direct_8x8_inference_flag:bool=False # u(1) frame_cropping_flag:bool=False # u(1) frame_crop_left_offset:int=0 # ue(v) frame_crop_right_offset:int=0 # ue(v) frame_crop_top_offset:int=0 # ue(v) frame_crop_bottom_offset:int=0 # ue(v) vui_parameters_present_flag:bool=False # u(1) vui:VUI|None=None # VUI object def __init__(self): self.scaling_list={} self.offset_for_ref_frame={} # TODO: ... # Compute options to pass to ffmpeg so as to reproduce the same SPS. # Very complex since some codec configuration are not provided by ffmpeg and/or libx264. # This is only an attempt for now and it is almost impossible to mimic any profile without # patching ffmpeg and/or libx264 to add the support for corner cases. def ffmpeg_options(self, video_id=0): x264opts = [] try: profile = {0x42:'baseline', 0x4D:'main', 0x64:'high', 0x6E:'high10', 0x7A:'high422', 0xF4:'high444'}[self.profile_idc] except KeyError: logger.error('Unknow profile: %x', self.profile_idc) return [] level = f'{floor(self.level_idc/10):d}.{self.level_idc % 10:d}' x264opts.extend([f'sps-id={self.seq_parameter_set_id:d}'] ) if self.bit_depth_chroma_minus8 not in [0,1,2,4,6,8]: logger.error('Bit depth of chrominance is not supported: %d', self.bit_depth_chroma_minus8+8) return [] if self.chroma_format_idc in range(4): if self.chroma_format_idc == 0: # Monochrome pass elif self.chroma_format_idc == 1: # YUV:4:2:0 pass elif self.chroma_format_idc == 2: # YUV:4:2:2 pass elif self.chroma_format_idc == 3: # YUV:4:4:4 pass else: logger.error('Unknow chrominance format: %x', self.chroma_format_idc) return [] res = [f'-profile:v:{video_id:d}', self.profile_idc, f'-level:v:{video_id:d}', level] return res def from_bytes(self, buf): logger.debug('Parsing: %s', hexdump.dump(buf,sep=':')) bit_position=0 # NAL Unit SPS bit_position, zero = read_bit(buf, bit_position) if zero != 0: raise ValueError(f'Reserved bit is not equal to 0: {zero:d}') bit_position, nal_ref_idc = read_bits(buf, bit_position,2) if nal_ref_idc != 3: raise ValueError(f'NAL ref idc is not equal to 3: {nal_ref_idc:d}') bit_position, nal_unit_type = read_bits(buf, bit_position,5) if nal_unit_type != 7: raise ValueError(f'NAL unit type is not a SPS: {nal_unit_type:d}') bit_position, self.profile_idc = read_byte(buf, bit_position) bit_position, self.constraint_set0_flag = read_bit(buf,bit_position) bit_position, self.constraint_set1_flag = read_bit(buf,bit_position) bit_position, self.constraint_set2_flag = read_bit(buf,bit_position) bit_position, self.constraint_set3_flag = read_bit(buf,bit_position) bit_position, self.constraint_set4_flag = read_bit(buf,bit_position) bit_position, self.constraint_set5_flag = read_bit(buf,bit_position) bit_position, v = read_bits(buf, bit_position, 2) if v!=0: raise ValueError(f'Reserved bits different from 0b00: {v:x}') bit_position, self.level_idc = read_byte(buf, bit_position) bit_position, self.seq_parameter_set_id = read_unsigned_exp_golomb(buf, bit_position) if self.profile_idc in [44, 83, 86, 100, 110, 118, 122, 128, 134, 135, 138, 139, 244]: bit_position, self.chroma_format_idc = read_unsigned_exp_golomb(buf, bit_position) if self.chroma_format_idc==3: bit_position, self.separate_colour_plane_flag=read_bit(buf, bit_position) bit_position, self.bit_depth_luma_minus8 = read_unsigned_exp_golomb(buf, bit_position) bit_position, self.bit_depth_chroma_minus8 = read_unsigned_exp_golomb(buf, bit_position) bit_position, self.qpprime_y_zero_transform_bypass_flag = read_boolean(buf, bit_position) bit_position, self.seq_scaling_matrix_present_flag = read_boolean(buf, bit_position) if self.seq_scaling_matrix_present_flag: nb_matrices = 12 if self.chroma_format_idc == 3 else 8 for i in range(nb_matrices): bit_position, present = read_boolean(buf, bit_position) if present: if i<6: bit_position, matrix = parse_scaling_list(buf, bit_position, 16) self.scaling_list[i] = matrix else: bit_position, matrix = parse_scaling_list(buf, bit_position, 64) self.scaling_list[i] = matrix else: self.scaling_list[i] = [] bit_position, self.log2_max_frame_num_minus4 = read_unsigned_exp_golomb(buf, bit_position) bit_position , self.pic_order_cnt_type = read_unsigned_exp_golomb(buf, bit_position) if self.pic_order_cnt_type == 0: bit_position, self.log2_max_pic_order_cnt_lsb_minus4 =\ read_unsigned_exp_golomb(buf, bit_position) elif self.pic_order_cnt_type == 1: bit_position, self.delta_pic_order_always_zero_flag = read_boolean(buf, bit_position) bit_position, self.offset_for_non_ref_pic = read_signed_exp_golomb(buf, bit_position) bit_position, self.offset_for_top_to_bottom_field = read_signed_exp_golomb(buf, bit_position) bit_position, self.num_ref_frames_in_pic_order_cnt_cycle =\ read_unsigned_exp_golomb(buf, bit_position) for i in range(self.num_ref_frames_in_pic_order_cnt_cycle): bit_position, v = read_unsigned_exp_golomb(buf, bit_position) self.offset_for_ref_frame[i]=v bit_position, self.max_num_ref_frames = read_unsigned_exp_golomb(buf, bit_position) bit_position, self.gaps_in_frame_num_value_allowed_flag = read_boolean(buf, bit_position) bit_position, self.pic_width_in_mbs_minus1 = read_unsigned_exp_golomb(buf, bit_position) bit_position, self.pic_height_in_map_units_minus1 = read_unsigned_exp_golomb(buf, bit_position) bit_position, self.frame_mbs_only_flag = read_boolean(buf, bit_position) if not self.frame_mbs_only_flag: bit_position, self.mb_adaptive_frame_field_flag = read_boolean(buf, bit_position) bit_position, self.direct_8x8_inference_flag = read_boolean(buf, bit_position) bit_position, self.frame_cropping_flag = read_boolean(buf, bit_position) if self.frame_cropping_flag: bit_position, self.frame_crop_left_offset = read_unsigned_exp_golomb(buf, bit_position) bit_position, self.frame_crop_right_offset = read_unsigned_exp_golomb(buf, bit_position) bit_position, self.frame_crop_top_offset = read_unsigned_exp_golomb(buf, bit_position) bit_position, self.frame_crop_bottom_offset = read_unsigned_exp_golomb(buf, bit_position) bit_position, self.vui_parameters_present_flag = read_boolean(buf, bit_position) if self.vui_parameters_present_flag: self.vui = VUI() bit_position = self.vui.from_bytes(buf,bit_position) logger.debug('VUI present: %s', self.vui) logger.debug('Parse end of SPS. Bit position: %d. Remaining bytes: %s.', bit_position, hexdump.dump(buf[floor(bit_position/8):], sep=':')) bit_position = parse_rbsp_trailing_bits(buf, bit_position) logger.debug('End of SPS: %d. Remaining bytes: %s', bit_position, hexdump.dump(buf[floor(bit_position/8):], sep=':')) return bit_position def to_bytes(self): buf = bytearray() bit_position = 0 bit_position = write_bit(buf, bit_position,0) bit_position = write_bits(buf, bit_position, 3, 2) bit_position = write_bits(buf, bit_position, 7, 5) bit_position = write_byte(buf, bit_position, self.profile_idc) bit_position = write_bit(buf, bit_position, self.constraint_set0_flag) bit_position = write_bit(buf, bit_position, self.constraint_set1_flag) bit_position = write_bit(buf, bit_position, self.constraint_set2_flag) bit_position = write_bit(buf, bit_position, self.constraint_set3_flag) bit_position = write_bit(buf, bit_position, self.constraint_set4_flag) bit_position = write_bit(buf, bit_position, self.constraint_set5_flag) bit_position = write_bits(buf, bit_position, 0, 2) bit_position = write_byte(buf, bit_position, self.level_idc) bit_position = write_unsigned_exp_golomb(buf, bit_position, self.seq_parameter_set_id) if self.profile_idc in [44, 83, 86, 100, 110, 118, 122, 128, 134, 135, 138, 139, 244]: bit_position = write_unsigned_exp_golomb(buf, bit_position, self.chroma_format_idc) if self.chroma_format_idc==3: bit_position = write_bit(buf, bit_position, self.separate_colour_plane_flag) bit_position = write_unsigned_exp_golomb(buf, bit_position, self.bit_depth_luma_minus8) bit_position = write_unsigned_exp_golomb(buf, bit_position, self.bit_depth_chroma_minus8) bit_position = write_boolean(buf, bit_position, self.qpprime_y_zero_transform_bypass_flag) bit_position = write_boolean(buf, bit_position, self.seq_scaling_matrix_present_flag) if self.seq_scaling_matrix_present_flag: nb_matrices = 12 if self.chroma_format_idc == 3 else 8 for i in range(nb_matrices): matrix = self.scaling_list[i] present = (len(matrix))!=0 bit_position = write_boolean(buf, bit_position, present) if present: if i<6: bit_position = write_scaling_list(buf, bit_position, 16, matrix) else: bit_position = write_scaling_list(buf, bit_position, 64, matrix) bit_position = write_unsigned_exp_golomb(buf, bit_position, self.log2_max_frame_num_minus4) bit_position = write_unsigned_exp_golomb(buf, bit_position, self.pic_order_cnt_type) if self.pic_order_cnt_type == 0: bit_position = write_unsigned_exp_golomb(buf, bit_position, self.log2_max_pic_order_cnt_lsb_minus4) elif self.pic_order_cnt_type == 1: bit_position = write_boolean(buf, bit_position, self.delta_pic_order_always_zero_flag) bit_position = write_signed_exp_golomb(buf, bit_position, self.offset_for_non_ref_pic) bit_position = write_signed_exp_golomb(buf, bit_position, self.offset_for_top_to_bottom_field) bit_position = write_unsigned_exp_golomb(buf, bit_position, self.num_ref_frames_in_pic_order_cnt_cycle) for i in range(self.num_ref_frames_in_pic_order_cnt_cycle): v = self.offset_for_ref_frame[i] bit_position = write_unsigned_exp_golomb(buf, bit_position, v) bit_position = write_unsigned_exp_golomb(buf, bit_position, self.max_num_ref_frames) bit_position = write_boolean(buf, bit_position, self.gaps_in_frame_num_value_allowed_flag) bit_position = write_unsigned_exp_golomb(buf, bit_position, self.pic_width_in_mbs_minus1) bit_position = write_unsigned_exp_golomb(buf, bit_position, self.pic_height_in_map_units_minus1) bit_position = write_boolean(buf, bit_position, self.frame_mbs_only_flag) if not self.frame_mbs_only_flag: bit_position = write_boolean(buf, bit_position, self.mb_adaptive_frame_field_flag) bit_position = write_boolean(buf, bit_position, self.direct_8x8_inference_flag) bit_position = write_boolean(buf, bit_position, self.frame_cropping_flag) if self.frame_cropping_flag: bit_position = write_unsigned_exp_golomb(buf, bit_position, self.frame_crop_left_offset) bit_position = write_unsigned_exp_golomb(buf, bit_position, self.frame_crop_right_offset) bit_position = write_unsigned_exp_golomb(buf, bit_position, self.frame_crop_top_offset) bit_position = write_unsigned_exp_golomb(buf, bit_position, self.frame_crop_bottom_offset) bit_position = write_boolean(buf, bit_position, self.vui_parameters_present_flag) if self.vui_parameters_present_flag: logger.debug('SPS has VUI. Writing VUI at position: %d', bit_position) bit_position = self.vui.to_bytes(buf, bit_position) logger.debug('VUI written. New bit position: %d', bit_position) bit_position = write_rbsp_trailing_bits(buf, bit_position) return buf @dataclass class PPS: pic_parameter_set_id:int=0 seq_parameter_set_id:int=0 entropy_coding_mode_flag:bool=False bottom_field_pic_order_in_frame_present_flag:bool=False num_slice_groups_minus1:int=0 slice_group_map_type:int=0 run_length_minus1:dict = field(default_factory=dict) top_left:dict = field(default_factory=dict) bottom_right:dict = field(default_factory=dict) slice_group_change_direction_flag:bool=False slice_group_change_rate_minus1:int=0 pic_size_in_map_units_minus1:int=0 slice_group_id:dict = field(default_factory=dict) num_ref_idx_l0_default_active_minus1:int=0 num_ref_idx_l2_default_active_minus1:int=0 weighted_pred_flag:bool=False weighted_bipred_idc:int=0 pic_init_qp_minus26:int=0 pic_init_qs_minus26:int=0 chroma_qp_index_offset:int=0 deblocking_filter_control_present_flag:bool=False constrained_intra_pred_flag:bool=False redundant_pic_cnt_present_flag:bool=False transform_8x8_mode_flag:bool=False pic_scaling_matrix_present_flag:bool=False pic_scaling_list:list[list[int]] = field(default_factory=list) second_chroma_qp_index_offset:int=0 def __init__(self): self.run_length_minus1={} self.top_left={} self.bottom_right={} self.slice_group_id={} self.pic_scaling_list=[] # PPS are located at byte boundary def from_bytes(self, buf, chroma_format_idc): logger.debug('Parsing: %s', (hexdump.dump(buf,sep=':'))) bit_position=0 # NAL Unit PPS bit_position, zero = read_bit(buf, bit_position) if zero != 0: raise ValueError(f'Reserved bit is not equal to 0: {zero:d}') bit_position, nal_ref_idc = read_bits(buf, bit_position,2) if nal_ref_idc != 3: raise ValueError(f'NAL ref idc is not equal to 3: {nal_ref_idc:d}') bit_position, nal_unit_type = read_bits(buf, bit_position,5) if nal_unit_type != 8: raise ValueError(f'NAL unit type is not a PPS: {nal_unit_type:d}') bit_position, self.pic_parameter_set_id = read_unsigned_exp_golomb(buf, bit_position) bit_position, self.seq_parameter_set_id = read_unsigned_exp_golomb(buf, bit_position) bit_position, self.entropy_coding_mode_flag = read_boolean(buf, bit_position) bit_position, self.bottom_field_pic_order_in_frame_present_flag =\ read_boolean(buf, bit_position) bit_position, self.num_slice_groups_minus1 = read_unsigned_exp_golomb(buf, bit_position) if self.num_slice_groups_minus1>0: bit_position, self.slice_group_map_type = read_unsigned_exp_golomb(buf, bit_position) if self.slice_group_map_type == 0: for i in range(self.num_slice_groups_minus1): bit_position, v = read_unsigned_exp_golomb(buf, bit_position) self.run_length_minus1[i]=v elif self.slice_group_map_type == 2: for i in range(self.num_slice_groups_minus1): bit_position, v = read_unsigned_exp_golomb(buf, bit_position) self.top_left[i] = v bit_position, v = read_unsigned_exp_golomb(buf, bit_position) self.bottom_right[i] = v elif self.slice_group_map_type in [3,4,5]: bit_position, self.slice_group_change_direction_flag = read_boolean(buf, bit_position) bit_position, self.slice_group_change_rate_minus1 =\ read_unsigned_exp_golomb(buf, bit_position) elif self.slice_group_map_type == 6: bit_position, self.pic_size_in_map_units_minus1 =\ read_unsigned_exp_golomb(buf, bit_position) nb_bits = ceil(log(self.num_slice_groups_minus1+1)) for i in range(self.pic_size_in_map_units_minus1): bit_position, v = read_bits(buf, bit_position, nb_bits) self.slice_group_id[i]=v bit_position, self.num_ref_idx_l0_default_active_minus1 =\ read_unsigned_exp_golomb(buf, bit_position) bit_position, self.num_ref_idx_l2_default_active_minus1 =\ read_unsigned_exp_golomb(buf, bit_position) bit_position, self.weighted_pred_flag = read_boolean(buf, bit_position) bit_position, self.weighted_bipred_idc = read_bits(buf, bit_position, 2) bit_position, self.pic_init_qp_minus26 = read_signed_exp_golomb(buf, bit_position) bit_position, self.pic_init_qs_minus26 = read_signed_exp_golomb(buf, bit_position) bit_position, self.chroma_qp_index_offset = read_signed_exp_golomb(buf, bit_position) bit_position, self.deblocking_filter_control_present_flag = read_boolean(buf, bit_position) bit_position, self.constrained_intra_pred_flag = read_boolean(buf, bit_position) bit_position, self.redundant_pic_cnt_present_flag = read_boolean(buf, bit_position) if more_rbsp_data(buf, bit_position): bit_position, self.transform_8x8_mode_flag = read_boolean(buf, bit_position) bit_position, self.pic_scaling_matrix_present_flag = read_boolean(buf, bit_position) if self.pic_scaling_matrix_present_flag: nb_matrices = 6 if chroma_format_idc == 3 else 2 if self.transform_8x8_mode_flag: nb_matrices+=6 else: nb_matrices = 6 for i in range(nb_matrices): bit_position, present = read_boolean(buf, bit_position) if present: if i<6: bit_position, matrix = parse_scaling_list(buf, bit_position, 16) self.pic_scaling_list.append(matrix) else: bit_position, matrix = parse_scaling_list(buf, bit_position, 64) self.pic_scaling_list.append(matrix) else: self.pic_scaling_list.append([]) bit_position, self.second_chroma_qp_index_offset = read_signed_exp_golomb(buf, bit_position) logger.info("parse RBSP") bit_position = parse_rbsp_trailing_bits(buf, bit_position) return bit_position def to_bytes(self, chroma_format_idc): buf = bytearray() bit_position = 0 # NAL Unit PPS bit_position = write_bit(buf, bit_position, 0) bit_position = write_bits(buf, bit_position, 3, 2) bit_position = write_bits(buf, bit_position, 8, 5) bit_position = write_unsigned_exp_golomb(buf, bit_position, self.pic_parameter_set_id) bit_position = write_unsigned_exp_golomb(buf, bit_position, self.seq_parameter_set_id) bit_position = write_boolean(buf, bit_position, self.entropy_coding_mode_flag) bit_position = write_boolean(buf, bit_position,\ self.bottom_field_pic_order_in_frame_present_flag) bit_position = write_unsigned_exp_golomb(buf, bit_position, self.num_slice_groups_minus1) if self.num_slice_groups_minus1>0: bit_position = write_unsigned_exp_golomb(buf, bit_position, self.slice_group_map_type) if self.slice_group_map_type == 0: for i in range(self.num_slice_groups_minus1): v = self.run_length_minus1[i] bit_position = write_unsigned_exp_golomb(buf, bit_position, v) elif self.slice_group_map_type == 2: for i in range(self.num_slice_groups_minus1): v = self.top_left[i] bit_position = write_unsigned_exp_golomb(buf, bit_position, v) v = self.bottom_right[i] bit_position = write_unsigned_exp_golomb(buf, bit_position, v) elif self.slice_group_map_type in [3,4,5]: bit_position = write_boolean(buf, bit_position, self.slice_group_change_direction_flag) bit_position = write_unsigned_exp_golomb(buf, bit_position, self.slice_group_change_rate_minus1) elif self.slice_group_map_type == 6: bit_position = write_unsigned_exp_golomb(buf, bit_position, self.pic_size_in_map_units_minus1) nb_bits = ceil(log(self.num_slice_groups_minus1+1)) for i in range(self.pic_size_in_map_units_minus1): v = self.slice_group_id[i] bit_position, v = write_bits(buf, bit_position, v, nb_bits) bit_position = write_unsigned_exp_golomb(buf, bit_position, self.num_ref_idx_l0_default_active_minus1) bit_position = write_unsigned_exp_golomb(buf, bit_position, self.num_ref_idx_l2_default_active_minus1) bit_position = write_boolean(buf, bit_position, self.weighted_pred_flag) bit_position = write_bits(buf, bit_position, self.weighted_bipred_idc, 2) bit_position = write_signed_exp_golomb(buf, bit_position, self.pic_init_qp_minus26) bit_position = write_signed_exp_golomb(buf, bit_position, self.pic_init_qs_minus26) bit_position = write_signed_exp_golomb(buf, bit_position, self.chroma_qp_index_offset) bit_position = write_boolean(buf, bit_position, self.deblocking_filter_control_present_flag) bit_position = write_boolean(buf, bit_position, self.constrained_intra_pred_flag) bit_position = write_boolean(buf, bit_position, self.redundant_pic_cnt_present_flag) bit_position = write_boolean(buf, bit_position, self.transform_8x8_mode_flag) bit_position = write_boolean(buf, bit_position, self.pic_scaling_matrix_present_flag) if self.pic_scaling_matrix_present_flag: nb_matrices = 6 if chroma_format_idc == 3 else 2 if self.transform_8x8_mode_flag: nb_matrices+=6 else: nb_matrices = 6 for i in range(nb_matrices): matrix = self.pic_scaling_list[i] logger.info("Retrieved pic scaling matrix: %s %d", matrix, len(matrix)) present = len(matrix)!=0 logger.info("Matrix is present: %s", present) bit_position = write_boolean(buf, bit_position, present) if present: if i<6: logger.info("Writing matrix: %s", matrix) bit_position = write_scaling_list(buf, bit_position, 16, matrix) else: logger.info("Writing matrix: %s", matrix) bit_position = write_scaling_list(buf, bit_position, 64, matrix) bit_position = write_signed_exp_golomb(buf, bit_position, self.second_chroma_qp_index_offset) bit_position = write_rbsp_trailing_bits(buf, bit_position) return buf