Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/selkies/webrtc/codecs/h265.py: 24%

1# 

2# This Source Code Form is subject to the terms of the Mozilla Public 

3# License, v. 2.0. If a copy of the MPL was not distributed with this 

4# file, You can obtain one at https://mozilla.org/MPL/2.0/. 

5# 

6"""RFC 7798 packetization of H.265 access units, the H.264 packer's shape 

7with the two-byte NAL header: one NAL per packet where it fits, an aggregation 

8packet (AP, type 48) for a run of small ones, fragmentation units (FU, type 49) 

9for a NAL larger than a packet. Every packet's payload header is a NAL header 

10whose type names the packet kind and whose layer id and TID come from the NALs 

11it carries.""" 

12 

13import math 

14from collections.abc import Iterable, Iterator 

15from struct import pack, unpack_from 

16from typing import Optional, Union 

17 

18from ..mediastreams import VIDEO_TIME_BASE, convert_timebase 

19from .base import Decoder, Encoder, EncodedPacket 

20from .h264 import PACKET_MAX, H264Encoder 

21 

22NAL_HEADER_SIZE = 2 

23NAL_TYPE_AP = 48 

24NAL_TYPE_FU = 49 

25FU_HEADER_SIZE = NAL_HEADER_SIZE + 1 

26LENGTH_FIELD_SIZE = 2 

27AP_HEADER_SIZE = NAL_HEADER_SIZE + LENGTH_FIELD_SIZE 

28# IRAP pictures (BLA, IDR, CRA, and the reserved IRAP types) open a key frame. 

29IRAP_NAL_TYPES = range(16, 24) 

30START_CODE = bytes([0, 0, 0, 1]) 

31 

32Buffer = Union[bytes, memoryview] 

33 

34 

35def nal_unit_type(nal: Buffer) -> int: 

36 return (nal[0] >> 1) & 0x3F 

37 

38 

39def h265_is_key(nals: Iterable[Buffer]) -> bool: 

40 return any(nal_unit_type(nal) in IRAP_NAL_TYPES for nal in nals) 

41 

42 

43class H265Decoder(Decoder): 

44 pass 

45 

46 

47class H265Encoder(Encoder): 

48 @staticmethod 

49 def _packetize_fu(data: Buffer) -> list[bytes]: 

50 available_size = PACKET_MAX - FU_HEADER_SIZE 

51 payload_size = len(data) - NAL_HEADER_SIZE 

52 num_packets = math.ceil(payload_size / available_size) 

53 num_larger_packets = payload_size % num_packets 

54 package_size = payload_size // num_packets 

55 

56 nal_type = nal_unit_type(data) 

57 # The payload header keeps F, the layer id, and the TID; its type is FU. 

58 payload_header = bytes([(data[0] & 0x81) | (NAL_TYPE_FU << 1), data[1]]) 

59 fu_header_start = payload_header + bytes([0x80 | nal_type]) 

60 fu_header_middle = payload_header + bytes([nal_type]) 

61 fu_header_end = payload_header + bytes([0x40 | nal_type]) 

62 fu_header = fu_header_start 

63 

64 packages = [] 

65 offset = NAL_HEADER_SIZE 

66 while offset < len(data): 

67 if num_larger_packets > 0: 

68 num_larger_packets -= 1 

69 payload = data[offset : offset + package_size + 1] 

70 offset += package_size + 1 

71 else: 

72 payload = data[offset : offset + package_size] 

73 offset += package_size 

74 if offset == len(data): 

75 fu_header = fu_header_end 

76 packages.append(b"".join((fu_header, payload))) 

77 fu_header = fu_header_middle 

78 assert offset == len(data), "incorrect fragment data" 

79 return packages 

80 

81 @staticmethod 

82 def _packetize_ap( 

83 data: Buffer, packages_iterator: Iterator[Buffer] 

84 ) -> tuple[bytes, Optional[Buffer]]: 

85 counter = 0 

86 available_size = PACKET_MAX - AP_HEADER_SIZE 

87 forbidden = 0 

88 layer_id = 0x3F 

89 tid = 0x07 

90 

91 payload = bytearray() 

92 try: 

93 nalu = data 

94 while len(nalu) <= available_size and counter < 9: 

95 forbidden |= nalu[0] & 0x80 

96 layer_id = min(layer_id, ((nalu[0] & 0x01) << 5) | (nalu[1] >> 3)) 

97 tid = min(tid, nalu[1] & 0x07) 

98 available_size -= LENGTH_FIELD_SIZE + len(nalu) 

99 counter += 1 

100 payload += pack("!H", len(nalu)) 

101 payload += nalu 

102 nalu = next(packages_iterator) 

103 if counter == 0: 

104 nalu = next(packages_iterator) 

105 except StopIteration: 

106 nalu = None 

107 

108 if counter <= 1: 

109 return (data if isinstance(data, bytes) else bytes(data)), nalu 

110 header = bytes([ 

111 forbidden | (NAL_TYPE_AP << 1) | (layer_id >> 5), 

112 ((layer_id & 0x1F) << 3) | tid, 

113 ]) 

114 return b"".join((header, payload)), nalu 

115 

116 @classmethod 

117 def _packetize(cls, packages: Iterable[Buffer]) -> list[bytes]: 

118 packetized_packages = [] 

119 packages_iterator = iter(packages) 

120 package = next(packages_iterator, None) 

121 while package is not None: 

122 if len(package) > PACKET_MAX: 

123 packetized_packages.extend(cls._packetize_fu(package)) 

124 package = next(packages_iterator, None) 

125 else: 

126 packetized, package = cls._packetize_ap(package, packages_iterator) 

127 packetized_packages.append(packetized) 

128 return packetized_packages 

129 

130 def pack(self, packet: EncodedPacket) -> tuple[list[bytes], int, bool]: 

131 nals = list(H264Encoder._split_bitstream(memoryview(packet.data))) 

132 timestamp = convert_timebase(packet.pts, packet.time_base, VIDEO_TIME_BASE) 

133 return self._packetize(nals), timestamp, h265_is_key(nals) 

134 

135 

136def h265_depayload(payload: bytes) -> bytes: 

137 """The Annex-B bytes one packet contributes to its access unit: whole NALs 

138 with start codes, a fragment's data with the NAL header rebuilt on the 

139 first fragment only.""" 

140 if len(payload) < NAL_HEADER_SIZE + 1: 

141 raise ValueError("NAL unit is too short") 

142 nal_type = nal_unit_type(payload) 

143 if nal_type < NAL_TYPE_AP: 

144 return START_CODE + payload 

145 if nal_type == NAL_TYPE_AP: 

146 output = bytearray() 

147 pos = NAL_HEADER_SIZE 

148 while pos < len(payload): 

149 if len(payload) < pos + LENGTH_FIELD_SIZE: 

150 raise ValueError("AP length field is truncated") 

151 size = unpack_from("!H", payload, pos)[0] 

152 pos += LENGTH_FIELD_SIZE 

153 if len(payload) < pos + size: 

154 raise ValueError("AP data is truncated") 

155 output += START_CODE 

156 output += payload[pos : pos + size] 

157 pos += size 

158 return bytes(output) 

159 if nal_type == NAL_TYPE_FU: 

160 fu_header = payload[NAL_HEADER_SIZE] 

161 data = payload[FU_HEADER_SIZE:] 

162 if fu_header & 0x80: 

163 nal_header = bytes([(payload[0] & 0x81) | ((fu_header & 0x3F) << 1), payload[1]]) 

164 return START_CODE + nal_header + data 

165 return data 

166 raise ValueError(f"NAL unit type {nal_type} is not supported")