3. H.264帧类型识别与NAL单元分类

任务目标

编写一个Python程序，读取H.264裸码流（Annex B格式），逐字节解析NAL单元，输出每个NAL单元的类型名称、负载大小及起始偏移，并进一步识别出I帧、P帧、B帧（即包含SPS/PPS/IDR/非IDR slice的NAL单元）。

环境依赖

仅需Python 3.6+标准库，无第三方依赖。输入文件为H.264 Annex B格式码流（以00 00 01或00 00 00 01作为起始码）。

关键完整代码

import sys
import os
from collections import namedtuple

NALUnit = namedtuple('NALUnit', ['offset', 'size', 'nal_type', 'frame_type'])

# H.264 NAL单元类型定义（表7-1）
NAL_TYPES = {
    1: '非IDR Slice (P/B帧)', 2: 'Slice分区A', 3: 'Slice分区B',
    4: 'Slice分区C', 5: 'IDR Slice (I帧)', 6: 'SEI', 7: 'SPS',
    8: 'PPS', 9: 'AUD', 10: '序列结束', 11: '码流结束', 12: '填充'
}

# slice_type映射（表7-6）：2,4,7,9 → I帧；0,3,5,8 → P帧；1,6 → B帧
SLICE_TYPE_MAP = {
    0: 'P帧', 1: 'B帧', 2: 'I帧', 3: 'P帧', 4: 'I帧',
    5: 'P帧', 6: 'B帧', 7: 'I帧', 8: 'P帧', 9: 'I帧'
}

def find_start_codes(data):
    """返回所有起始码的偏移位置（含3字节和4字节起始码）"""
    positions = []
    i = 0
    while i < len(data) - 3:
        if data[i:i+3] == b'\x00\x00\x01':
            positions.append(i)
            i += 3
        elif data[i:i+4] == b'\x00\x00\x00\x01':
            positions.append(i)
            i += 4
        else:
            i += 1
    return positions

def parse_nal_header(byte):
    """解析NAL头： forbidden_zero_bit(1) | nal_ref_idc(2) | nal_unit_type(5)"""
    nal_type = byte & 0x1F
    nal_ref_idc = (byte & 0x60) >> 5
    return nal_type, nal_ref_idc

def get_frame_type(nal_payload, nal_type):
    """从slice头中提取slice_type，判断帧类型"""
    if nal_type not in (1, 5):
        return None
    # 解析slice_header：first_mb_in_slice使用ue(v)，slice_type使用ue(v)
    # 简化实现：跳过first_mb_in_slice的Exp-Golomb编码
    i = 0
    # 读取第一个ue(v) —— first_mb_in_slice
    leading_zeros = 0
    while i < len(nal_payload) and nal_payload[i] == 0:
        leading_zeros += 1
        i += 1
    if i >= len(nal_payload):
        return None
    # 跳过first_mb_in_slice的剩余部分
    code_num = (1 << leading_zeros) - 1 + nal_payload[i]
    i += 1
    # 跳过first_mb_in_slice的额外信息位
    bits_to_skip = leading_zeros
    while bits_to_skip > 0:
        byte_idx = i + (bits_to_skip - 1) // 8
        if byte_idx >= len(nal_payload):
            return None
        bits_to_skip -= 1
    i = byte_idx + 1
    
    # 读取第二个ue(v) —— slice_type
    if i >= len(nal_payload):
        return None
    leading_zeros = 0
    while i < len(nal_payload) and nal_payload[i] == 0:
        leading_zeros += 1
        i += 1
    if i >= len(nal_payload):
        return None
    slice_type = (1 << leading_zeros) - 1 + nal_payload[i]
    return SLICE_TYPE_MAP.get(slice_type, '未知')

def analyze_h264(filepath):
    with open(filepath, 'rb') as f:
        data = f.read()
    
    start_positions = find_start_codes(data)
    nal_units = []
    
    for idx, pos in enumerate(start_positions):
        # 确定NAL单元结束位置（下一个起始码或文件末尾）
        if idx + 1 < len(start_positions):
            end = start_positions[idx + 1]
        else:
            end = len(data)
        
        # 跳过起始码
        if data[pos:pos+4] == b'\x00\x00\x00\x01':
            nal_start = pos + 4
        else:
            nal_start = pos + 3
        
        if nal_start >= end:
            continue
        
        nal_type, nal_ref_idc = parse_nal_header(data[nal_start])
        payload = data[nal_start+1:end]  # 去掉NAL头后的负载
        
        frame_type = get_frame_type(payload, nal_type)
        nal_units.append(NALUnit(
            offset=pos,
            size=end - pos,
            nal_type=nal_type,
            frame_type=frame_type
        ))
    
    return nal_units

def main():
    if len(sys.argv) != 2:
        print('用法: python h264_nal_parser.py <输入文件.h264>')
        sys.exit(1)
    
    filepath = sys.argv[1]
    if not os.path.exists(filepath):
        print(f'错误: 文件不存在 - {filepath}')
        sys.exit(1)
    
    nal_units = analyze_h264(filepath)
    
    print(f'文件: {filepath}')
    print(f'总大小: {os.path.getsize(filepath)} 字节')
    print(f'NAL单元总数: {len(nal_units)}\n')
    print(f'{"偏移":>10} {"大小":>8} {"类型":>5} {"描述":>25} {"帧类型":>10}')
    print('-' * 65)
    
    for nal in nal_units:
        desc = NAL_TYPES.get(nal.nal_type, f'未知({nal.nal_type})')
        frame = nal.frame_type if nal.frame_type else '-'
        print(f'{nal.offset:>10} {nal.size:>8} {nal.nal_type:>5} {desc:>25} {frame:>10}')

if __name__ == '__main__':
    main()

运行方式

python h264_nal_parser.py input.h264

验证方法

使用已知结构的H.264测试码流验证。例如，用ffmpeg生成测试文件：

ffmpeg -f lavfi -i testsrc=duration=1:size=320x240:rate=30 -c:v libx264 -g 15 test.h264

运行解析器后，预期输出应包含：

• 至少一个SPS（类型7）和PPS（类型8）NAL单元，位于文件开头附近；
• 首个图像帧为IDR Slice（类型5，帧类型标记为“I帧”）；
• 后续帧为类型1（非IDR Slice），帧类型为“P帧”或“B帧”（取决于编码配置）；
• 所有NAL单元的偏移与大小之和等于文件总大小。

判定通过条件：程序能正确输出每个NAL的类型编号与名称，且I/P/B帧标记与用ffprobe查看的帧类型一致（ffprobe -show_frames test.h264）。

常见问题

• slice_type解析失败：若码流包含多个slice或使用了FMO（灵活宏块排序），本简化解析器可能无法正确读取slice_type。此时可改用ffprobe交叉验证帧类型。
• 字节流包含00 00 01：某些码流内部可能包含伪起始码，但Annex B格式中起始码后紧跟的字节若为0x00或0x01则可能是伪起始码，本程序未处理此情况，对标准编码器输出通常无影响。