Separate H264 specialized functions.

src/tscut/h264/bitstream.py

@@ -0,0 +1,300 @@
+# SPDX-License-Identifier: GPL-2.0-or-later
+#
+# Copyright (C) 2026 Frédéric Tronel
+
+""" All the following code is a transposition of documents:
+    ISO/IEC H.264-201602
+    ISO/IEC 14496-15
+"""
+
+import logging
+from math import floor, log
+
+from typeguard import typechecked
+import hexdump
+
+
+logger = logging.getLogger(__name__)
+
+@typechecked
+def read_bit(buf:bytes, bit_position: int) -> tuple[int, int]:
+    """
+    Read a single bit from a byte buffer.
+
+    This function is part of the implementation of the H.264/AVC video compression standard,
+    as specified in ISO/IEC H.264-201602 and ISO/IEC 14496-15.
+    It takes a byte buffer and a bit position as input, and returns a tuple containing
+    the updated bit position and the value of the bit at the specified position.
+
+    Args:
+        buf (bytes): The byte buffer to read from.
+        bit_position (int): The position of the bit to read, starting from 0.
+
+    Returns:
+        tuple[int, int]: A tuple containing the updated bit position (bit_position + 1) and the
+        value of the bit (0 or 1).
+    """
+    byte_position = floor(floor(bit_position/8))
+    byte = buf[byte_position]
+    bit = (byte >> (7-(bit_position % 8))) & 1
+    return bit_position+1, bit
+
+@typechecked
+def read_boolean(buf:bytes, bit_position: int) -> tuple[int, bool]:
+    """
+    Read a boolean value from a byte buffer.
+
+    This function reads a single bit from the byte buffer at the specified position and interprets
+    it as a boolean value.
+    It returns a tuple containing the updated bit position and the boolean value.
+
+    Args:
+        buf (bytes): The byte buffer to read from.
+        bit_position (int): The position of the bit to read, starting from 0.
+
+    Returns:
+        tuple[int, bool]: A tuple containing the updated bit position and the boolean value 
+        (True if the bit is 1, False if the bit is 0).
+    """
+    bit_position, b = read_bit(buf, bit_position)
+    return bit_position, b==1
+
+@typechecked
+def read_bits(buf:bytes, bit_position: int, nb_bits: int)  -> tuple[int, int]:
+    v = 0
+    for _ in range(nb_bits):
+        bit_position, bit = read_bit(buf, bit_position)
+        v = v*2+bit
+    return bit_position, v
+
+@typechecked
+def read_byte(buf:bytes, bit_position: int)  -> tuple[int, int]:
+    bit_position, b = read_bits(buf, bit_position, 8)
+    return bit_position, b
+
+@typechecked
+def read_word(buf:bytes, bit_position: int) -> tuple[int, int]:
+    bit_position, w = read_bits(buf, bit_position, 16)
+    return bit_position, w
+
+@typechecked
+def read_long(buf:bytes, bit_position: int) -> tuple[int, int]:
+    bit_position, value = read_bits(buf, bit_position, 32)
+    return bit_position, value
+
+@typechecked
+def read_unsigned_exp_golomb(buf:bytes, bit_position: int) -> tuple[int, int]:
+    nb_zeroes=0
+    while True:
+        bit_position, b = read_bit(buf, bit_position)
+        if b!=0:
+            break
+        nb_zeroes+=1
+
+    v1 = 1
+    bit_position, v2 = read_bits(buf, bit_position, nb_zeroes)
+    v = (v1<<nb_zeroes)+v2
+    return bit_position, v-1
+
+@typechecked
+def read_signed_exp_golomb(buf:bytes, bit_position: int) -> tuple[int, int]:
+    bit_position, v = read_unsigned_exp_golomb(buf, bit_position)
+    match v%2:
+        case 0:
+            return bit_position, -(v>>1)
+        case 1:
+            return bit_position, (v+1)>>1
+
+@typechecked
+def write_bit(buf:bytes, bit_position: int, b) -> int:
+    buf_length = len(buf)
+    byte_position = floor(bit_position/8)
+
+    if byte_position >= buf_length:
+        extension = bytearray(byte_position+1-buf_length)
+        buf.extend(extension)
+
+    buf[byte_position] |= (b<<(7-(bit_position % 8)))
+    bit_position+=1
+
+    return bit_position
+
+@typechecked
+def write_boolean(buf:bytes, bit_position: int, b: bool) -> int:
+    if b:
+        bit_position = write_bit(buf, bit_position, 1)
+    else:
+        bit_position = write_bit(buf, bit_position, 0)
+    return bit_position
+
+@typechecked
+def write_bits(buf:bytes, bit_position: int, v, size) -> int:
+    for i in range(size-1,-1,-1):
+        b = (v>>i)&1
+        bit_position = write_bit(buf, bit_position, b)
+
+    return bit_position
+
+@typechecked
+def write_byte(buf:bytes, bit_position: int, v) -> int:
+    bit_position = write_bits(buf, bit_position, v, 8)
+    return bit_position
+
+@typechecked
+def write_word(buf:bytes, bit_position: int, v) -> int:
+    bit_position = write_bits(buf, bit_position, v, 16)
+    return bit_position
+
+@typechecked
+def write_long(buf:bytes, bit_position: int, v) -> int:
+    bit_position = write_bits(buf, bit_position, v, 32)
+    return bit_position
+
+@typechecked
+def write_unsigned_exp_golomb(buf:bytes, bit_position: int, v) -> int:
+    n = floor(log(v+1)/log(2))+1
+    # Write zeroes
+    bit_position = write_bits(buf, bit_position, 0, n-1)
+    bit_position = write_bit(buf, bit_position, 1)
+    bit_position = write_bits(buf, bit_position, v+1, n-1)
+
+    return bit_position
+
+@typechecked
+def write_signed_exp_golomb(buf:bytes, bit_position: int, v) -> int:
+    if v <= 0:
+        bit_position = write_unsigned_exp_golomb(buf, bit_position, -v*2)
+    else:
+        bit_position = write_unsigned_exp_golomb(buf, bit_position, v*2-1)
+
+    return bit_position
+
+@typechecked
+def parse_rbsp_trailing_bits(buf:bytes, bit_position: int) -> int:
+    bit_position, one = read_bit(buf, bit_position)
+    if one==0:
+        raise ValueError(f'Stop bit should be equal to one. Read: {one:d}')
+    while bit_position%8 != 0:
+        bit_position, zero = read_bit(buf, bit_position)
+        if zero==1:
+            raise ValueError('Trailing bit should be equal to zero')
+
+    return bit_position
+
+@typechecked
+def write_rbsp_trailing_bits(buf:bytes, bit_position: int) -> int:
+    bit_position = write_bit(buf, bit_position, 1)
+    while bit_position%8 != 0:
+        bit_position = write_bit(buf, bit_position, 0)
+
+    return bit_position
+
+@typechecked
+def more_rbsp_data(buf:bytes, bit_position: int) -> bool:
+    logger.debug('Is there more data in buffer of length: %d at bit position: %d',
+                 len(buf), bit_position)
+
+    byte_length = len(buf)
+    bit_length = byte_length*8
+
+    # We are at the end of buffer
+    if bit_position == bit_length:
+        return False
+
+    found = False
+    for i in range(bit_length-1,-1,-1):
+        pos, b = read_bit(buf, i)
+        if b == 1:
+            found = True
+            break
+
+    if not found:
+        raise ValueError('Impossible to find trailing stop bit !')
+
+    # No more data
+    if bit_position == pos:
+        return False
+
+    return True
+
+# Convert from RBSP (Raw Byte Sequence Payload)  to SODB (String Of Data Bits)
+@typechecked
+def rbsp_to_sodb(buf:bytes) -> bytes:
+    logger.debug('RBSP: %s', hexdump.dump(buf, sep=':'))
+
+    res = buf
+    for b in [ b'\x00', b'\x01', b'\x02', b'\x03']:
+        pattern = b'\x00\x00\x03'+b
+        replacement = b'\x00\x00' + b
+        res = res.replace(pattern, replacement)
+
+    logger.debug('SODB: %s', hexdump.dump(res, sep=':'))
+    return res
+
+# Reverse operation SODB to RBSP.
+@typechecked
+def sodb_to_rbsp(buf:bytes) -> bytes:
+    logger.debug('SODB: %s', hexdump.dump(buf, sep=':'))
+
+    res = buf
+    for b in [ b'\x03', b'\x00', b'\x01', b'\x02']:
+        pattern = b'\x00\x00'+b
+        replacement = b'\x00\x00\x03' + b
+        res = res.replace(pattern, replacement)
+
+    logger.debug('RBSP: %s', hexdump.dump(res, sep=':'))
+    return res
+
+# Useful for SPS and PPS
+@typechecked
+def parse_scaling_list(buf:bytes, bit_position: int, size) -> tuple[int,list[int]]:
+    res = []
+    last_scale = 8
+    next_scale = 8
+    for _ in range(size):
+        if next_scale != 0:
+            bit_position, delta_scale = read_signed_exp_golomb(buf, bit_position)
+            next_scale = (last_scale+delta_scale+256) % 256
+        v = last_scale if next_scale==0 else next_scale
+        res.append(v)
+        last_scale = v
+
+    return bit_position,res
+
+# TODO: test optimized version.
+# The ISO/IEC H.264-201602 seems to take into account the case where the end of the deltas list
+#  is full of zeroes.
+@typechecked
+def write_scaling_list(buf:bytes, bit_position: int, size, matrix:list[int],
+                       optimized: bool = False) -> int:
+    logger.debug('Dumping matrix: %s of size: %d, size parameter: %d.', matrix, len(matrix), size)
+
+    prev = 8
+    deltas = []
+    for i in range(size):
+        v = matrix[i]
+        delta = v - prev
+        deltas.append(delta)
+        prev = v
+
+    if not optimized:
+        for delta in deltas:
+            bit_position = write_signed_exp_golomb(buf, bit_position, delta)
+    else:
+        logger.error('Not yet implemented')
+        exit(-1)
+        # reverse = deltas.reverse()
+        # compressed = False
+        # while len(reverse)>0:
+        #     if reverse[0] == 0:
+        #         compressed = True
+        #         reverse.pop()
+        #     else:
+        #         break
+        # deltas = reverse.reverse()
+        # if compressed:
+        #     deltas.append(0)
+        # for delta in deltas:
+        #     bit_position = write_signed_exp_golomb(buf, bit_position, delta)
+
+    return bit_position

src/tscut/tscut.py

@@ -34,6 +34,32 @@ from iso639.exceptions import InvalidLanguageValue
 from tqdm import tqdm
 from typeguard import typechecked
 
+from tscut.h264.bitstream import (
+    read_bit,
+    read_bits,
+    read_boolean,
+    read_byte,
+    read_word,
+    read_long,
+    read_unsigned_exp_golomb,
+    read_signed_exp_golomb,
+    write_bit,
+    write_boolean,
+    write_bits,
+    write_byte,
+    write_word,
+    write_long,
+    write_unsigned_exp_golomb,
+    write_signed_exp_golomb,
+    parse_scaling_list,
+    write_scaling_list,
+    write_rbsp_trailing_bits,
+    parse_rbsp_trailing_bits,
+    sodb_to_rbsp,
+    rbsp_to_sodb,
+    more_rbsp_data
+)
+
 # Local modules
 # TODO: create local modules for MP4, MKV
 
@@ -218,7 +244,6 @@ def get_frame_rate(ffprobe_path:str, input_file: IO[bytes]) -> float|None:
             return None
         if abs(frame_rate1*2 - frame_rate2) < 0.2:
             return frame_rate2/2
-        else:
         logger.error('Video is interlaced and the disperancy between frame rates is too big:\
             %f / %f', frame_rate1, frame_rate2)
         return None
@@ -443,296 +468,7 @@ def get_codec_private_data_from_mkv(mkvinfo_path:str,
     return None, None
 
 
-# All the following code is a transposition of documents:
-# ISO/IEC H.264-201602
-# ISO/IEC 14496-15
-@typechecked
-def read_bit(buf:bytes, bit_position: int) -> tuple[int, int]:
-    """
-    Read a single bit from a byte buffer.
 
-    This function is part of the implementation of the H.264/AVC video compression standard,
-    as specified in ISO/IEC H.264-201602 and ISO/IEC 14496-15.
-    It takes a byte buffer and a bit position as input, and returns a tuple containing
-    the updated bit position and the value of the bit at the specified position.
-
-    Args:
-        buf (bytes): The byte buffer to read from.
-        bit_position (int): The position of the bit to read, starting from 0.
-
-    Returns:
-        tuple[int, int]: A tuple containing the updated bit position (bit_position + 1) and the
-        value of the bit (0 or 1).
-    """
-    byte_position = floor(floor(bit_position/8))
-    byte = buf[byte_position]
-    bit = (byte >> (7-(bit_position % 8))) & 1
-    return bit_position+1, bit
-
-@typechecked
-def read_boolean(buf:bytes, bit_position: int) -> tuple[int, bool]:
-    """
-    Read a boolean value from a byte buffer.
-
-    This function reads a single bit from the byte buffer at the specified position and interprets
-    it as a boolean value.
-    It returns a tuple containing the updated bit position and the boolean value.
-
-    Args:
-        buf (bytes): The byte buffer to read from.
-        bit_position (int): The position of the bit to read, starting from 0.
-
-    Returns:
-        tuple[int, bool]: A tuple containing the updated bit position and the boolean value 
-        (True if the bit is 1, False if the bit is 0).
-    """
-    bit_position, b = read_bit(buf, bit_position)
-    return bit_position, b==1
-
-@typechecked
-def read_bits(buf:bytes, bit_position: int, nb_bits: int)  -> tuple[int, int]:
-    v = 0
-    for _ in range(nb_bits):
-        bit_position, bit = read_bit(buf, bit_position)
-        v = v*2+bit
-    return bit_position, v
-
-@typechecked
-def read_byte(buf:bytes, bit_position: int)  -> tuple[int, int]:
-    bit_position, b = read_bits(buf, bit_position, 8)
-    return bit_position, b
-
-@typechecked
-def read_word(buf:bytes, bit_position: int) -> tuple[int, int]:
-    bit_position, w = read_bits(buf, bit_position, 16)
-    return bit_position, w
-
-@typechecked
-def read_long(buf:bytes, bit_position: int) -> tuple[int, int]:
-    bit_position, value = read_bits(buf, bit_position, 32)
-    return bit_position, value
-
-@typechecked
-def read_unsigned_exp_golomb(buf:bytes, bit_position: int) -> tuple[int, int]:
-    nb_zeroes=0
-    while True:
-        bit_position, b = read_bit(buf, bit_position)
-        if b!=0:
-            break
-        nb_zeroes+=1
-
-    v1 = 1
-    bit_position, v2 = read_bits(buf, bit_position, nb_zeroes)
-    v = (v1<<nb_zeroes)+v2
-    return bit_position, v-1
-
-@typechecked
-def read_signed_exp_golomb(buf:bytes, bit_position: int) -> tuple[int, int]:
-    bit_position, v = read_unsigned_exp_golomb(buf, bit_position)
-    match v%2:
-        case 0:
-            return bit_position, -(v>>1)
-        case 1:
-            return bit_position, (v+1)>>1
-
-@typechecked
-def write_bit(buf:bytes, bit_position: int, b) -> int:
-    buf_length = len(buf)
-    byte_position = floor(bit_position/8)
-
-    if byte_position >= buf_length:
-        extension = bytearray(byte_position+1-buf_length)
-        buf.extend(extension)
-
-    buf[byte_position] |= (b<<(7-(bit_position % 8)))
-    bit_position+=1
-
-    return bit_position
-
-@typechecked
-def write_boolean(buf:bytes, bit_position: int, b: bool) -> int:
-    if b:
-        bit_position = write_bit(buf, bit_position, 1)
-    else:
-        bit_position = write_bit(buf, bit_position, 0)
-    return bit_position
-
-@typechecked
-def write_bits(buf:bytes, bit_position: int, v, size) -> int:
-    for i in range(size-1,-1,-1):
-        b = (v>>i)&1
-        bit_position = write_bit(buf, bit_position, b)
-
-    return bit_position
-
-@typechecked
-def write_byte(buf:bytes, bit_position: int, v) -> int:
-    bit_position = write_bits(buf, bit_position, v, 8)
-    return bit_position
-
-@typechecked
-def write_word(buf:bytes, bit_position: int, v) -> int:
-    bit_position = write_bits(buf, bit_position, v, 16)
-    return bit_position
-
-@typechecked
-def write_long(buf:bytes, bit_position: int, v) -> int:
-    bit_position = write_bits(buf, bit_position, v, 32)
-    return bit_position
-
-@typechecked
-def write_unsigned_exp_golomb(buf:bytes, bit_position: int, v) -> int:
-    n = floor(log(v+1)/log(2))+1
-    # Write zeroes
-    bit_position = write_bits(buf, bit_position, 0, n-1)
-    bit_position = write_bit(buf, bit_position, 1)
-    bit_position = write_bits(buf, bit_position, v+1, n-1)
-
-    return bit_position
-
-@typechecked
-def write_signed_exp_golomb(buf:bytes, bit_position: int, v) -> int:
-    if v <= 0:
-        bit_position = write_unsigned_exp_golomb(buf, bit_position, -v*2)
-    else:
-        bit_position = write_unsigned_exp_golomb(buf, bit_position, v*2-1)
-
-    return bit_position
-
-@typechecked
-def parse_rbsp_trailing_bits(buf:bytes, bit_position: int) -> int:
-    bit_position, one = read_bit(buf, bit_position)
-    if one==0:
-        raise ValueError(f'Stop bit should be equal to one. Read: {one:d}')
-    while bit_position%8 != 0:
-        bit_position, zero = read_bit(buf, bit_position)
-        if zero==1:
-            raise ValueError('Trailing bit should be equal to zero')
-
-    return bit_position
-
-@typechecked
-def write_rbsp_trailing_bits(buf:bytes, bit_position: int) -> int:
-    bit_position = write_bit(buf, bit_position, 1)
-    while bit_position%8 != 0:
-        bit_position = write_bit(buf, bit_position, 0)
-
-    return bit_position
-
-@typechecked
-def more_rbsp_data(buf:bytes, bit_position: int) -> bool:
-    logger = logging.getLogger(__name__)
-    logger.debug('Is there more data in buffer of length: %d at bit position: %d',
-                 len(buf), bit_position)
-
-    byte_length = len(buf)
-    bit_length = byte_length*8
-
-    # We are at the end of buffer
-    if bit_position == bit_length:
-        return False
-    else:
-        found = False
-        for i in range(bit_length-1,-1,-1):
-            pos, b = read_bit(buf, i)
-            if b == 1:
-                found = True
-                break
-
-    if not found:
-        raise ValueError('Impossible to find trailing stop bit !')
-
-    # No more data
-    if bit_position == pos:
-        return False
-
-    return True
-
-# Convert from RBSP (Raw Byte Sequence Payload)  to SODB (String Of Data Bits)
-@typechecked
-def rbsp_to_sodb(buf:bytes) -> bytes:
-    logger = logging.getLogger(__name__)
-
-    logger.debug('RBSP: %s', hexdump.dump(buf, sep=':'))
-
-    res = buf
-    for b in [ b'\x00', b'\x01', b'\x02', b'\x03']:
-        pattern = b'\x00\x00\x03'+b
-        replacement = b'\x00\x00' + b
-        res = res.replace(pattern, replacement)
-
-    logger.debug('SODB: %s', hexdump.dump(res, sep=':'))
-    return res
-
-# Reverse operation SODB to RBSP.
-@typechecked
-def sodb_to_rbsp(buf:bytes) -> bytes:
-    logger = logging.getLogger(__name__)
-    logger.debug('SODB: %s', hexdump.dump(buf, sep=':'))
-
-    res = buf
-    for b in [ b'\x03', b'\x00', b'\x01', b'\x02']:
-        pattern = b'\x00\x00'+b
-        replacement = b'\x00\x00\x03' + b
-        res = res.replace(pattern, replacement)
-
-    logger.debug('RBSP: %s', hexdump.dump(res, sep=':'))
-    return res
-
-# Useful for SPS and PPS
-@typechecked
-def parse_scaling_list(buf:bytes, bit_position: int, size) -> tuple[int,list[int]]:
-    res = []
-    last_scale = 8
-    next_scale = 8
-    for _ in range(size):
-        if next_scale != 0:
-            bit_position, delta_scale = read_signed_exp_golomb(buf, bit_position)
-            next_scale = (last_scale+delta_scale+256) % 256
-        v = last_scale if next_scale==0 else next_scale
-        res.append(v)
-        last_scale = v
-
-    return bit_position,res
-
-# TODO: test optimized version.
-# The ISO/IEC H.264-201602 seems to take into account the case where the end of the deltas list
-#  is full of zeroes.
-@typechecked
-def write_scaling_list(buf:bytes, bit_position: int, size, matrix:list[int],
-                       optimized: bool = False) -> int:
-    logger = logging.getLogger(__name__)
-    logger.debug('Dumping matrix: %s of size: %d, size parameter: %d.', matrix, len(matrix), size)
-
-    prev = 8
-    deltas = []
-    for i in range(size):
-        v = matrix[i]
-        delta = v - prev
-        deltas.append(delta)
-        prev = v
-
-    if not optimized:
-        for delta in deltas:
-            bit_position = write_signed_exp_golomb(buf, bit_position, delta)
-    else:
-        logger.error('Not yet implemented')
-        exit(-1)
-        # reverse = deltas.reverse()
-        # compressed = False
-        # while len(reverse)>0:
-        #     if reverse[0] == 0:
-        #         compressed = True
-        #         reverse.pop()
-        #     else:
-        #         break
-        # deltas = reverse.reverse()
-        # if compressed:
-        #     deltas.append(0)
-        # for delta in deltas:
-        #     bit_position = write_signed_exp_golomb(buf, bit_position, delta)
-
-    return bit_position
 
 @dataclass
 class HRD: