717 lines
39 KiB
Python
717 lines
39 KiB
Python
# SPDX-License-Identifier: GPL-2.0-or-later
|
|
#
|
|
# Copyright (C) 2026 Frédéric Tronel
|
|
|
|
import logging
|
|
from dataclasses import dataclass, field
|
|
from math import ceil, floor, log
|
|
|
|
import hexdump
|
|
|
|
from tscut.h264.bitstream import (
|
|
more_rbsp_data,
|
|
parse_rbsp_trailing_bits,
|
|
parse_scaling_list,
|
|
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,
|
|
)
|
|
|
|
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
|