Shortcuts on this page
r m x toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1# This Source Code Form is subject to the terms of the Mozilla Public
2# License, v. 2.0. If a copy of the MPL was not distributed with this
3# file, You can obtain one at https://mozilla.org/MPL/2.0/.
4#
5# This file incorporates work covered by the following copyright and
6# permission notice:
7#
8# Copyright (c) Jeremy Lainé.
9# All rights reserved.
10#
11# Redistribution and use in source and binary forms, with or without
12# modification, are permitted provided that the following conditions are met:
13#
14# * Redistributions of source code must retain the above copyright notice,
15# this list of conditions and the following disclaimer.
16# * Redistributions in binary form must reproduce the above copyright notice,
17# this list of conditions and the following disclaimer in the documentation
18# and/or other materials provided with the distribution.
19# * Neither the name of aiortc nor the names of its contributors may
20# be used to endorse or promote products derived from this software without
21# specific prior written permission.
22#
23# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
24# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
25# WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
26# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
27# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
28# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
29# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
30# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
31# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
32# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
34import logging
35import math
36import re
37from collections.abc import Iterable, Iterator, Sequence
38from itertools import tee
39from struct import pack, unpack_from
40from typing import Optional, Type, TypeVar, Union
42from ..mediastreams import VIDEO_TIME_BASE, convert_timebase
43from .base import Decoder, Encoder, EncodedPacket
45logger = logging.getLogger(__name__)
47DEFAULT_BITRATE = 1000000 # 1 Mbps
48MIN_BITRATE = 500000 # 500 kbps
49MAX_BITRATE = 3000000 # 3 Mbps
51MAX_FRAME_RATE = 30
52PACKET_MAX = 1200
54NAL_TYPE_FU_A = 28
55NAL_TYPE_STAP_A = 24
57NAL_HEADER_SIZE = 1
58FU_A_HEADER_SIZE = 2
59LENGTH_FIELD_SIZE = 2
60STAP_A_HEADER_SIZE = NAL_HEADER_SIZE + LENGTH_FIELD_SIZE
62# Annex-B start-code scanner; compiled once. re searches run at C speed over a
63# memoryview, which bytes.find cannot do without materializing the buffer.
64NAL_START_CODE = re.compile(b"\x00\x00\x01")
66DESCRIPTOR_T = TypeVar("DESCRIPTOR_T", bound="H264PayloadDescriptor")
67T = TypeVar("T")
70def pairwise(iterable: Sequence[T]) -> Iterator[tuple[T, T]]:
71 a, b = tee(iterable)
72 next(b, None)
73 return zip(a, b)
76class H264PayloadDescriptor:
77 def __init__(self, first_fragment: bool) -> None:
78 self.first_fragment = first_fragment
80 def __repr__(self) -> str:
81 return f"H264PayloadDescriptor(FF={self.first_fragment})"
83 @classmethod
84 def parse(cls: Type[DESCRIPTOR_T], data: bytes) -> tuple[DESCRIPTOR_T, bytes]:
85 output = bytes()
87 # NAL unit header
88 if len(data) < 2:
89 raise ValueError("NAL unit is too short")
90 nal_type = data[0] & 0x1F
91 f_nri = data[0] & (0x80 | 0x60)
92 pos = NAL_HEADER_SIZE
94 if nal_type in range(1, 24):
95 # single NAL unit
96 output = bytes([0, 0, 0, 1]) + data
97 obj = cls(first_fragment=True)
98 elif nal_type == NAL_TYPE_FU_A:
99 # fragmentation unit
100 original_nal_type = data[pos] & 0x1F
101 first_fragment = bool(data[pos] & 0x80)
102 pos += 1
104 if first_fragment:
105 original_nal_header = bytes([f_nri | original_nal_type])
106 output += bytes([0, 0, 0, 1])
107 output += original_nal_header
108 output += data[pos:]
110 obj = cls(first_fragment=first_fragment)
111 elif nal_type == NAL_TYPE_STAP_A:
112 # single time aggregation packet
113 offsets = []
114 while pos < len(data):
115 if len(data) < pos + LENGTH_FIELD_SIZE:
116 raise ValueError("STAP-A length field is truncated")
117 nalu_size = unpack_from("!H", data, pos)[0]
118 pos += LENGTH_FIELD_SIZE
119 offsets.append(pos)
121 pos += nalu_size
122 if len(data) < pos:
123 raise ValueError("STAP-A data is truncated")
125 offsets.append(len(data) + LENGTH_FIELD_SIZE)
126 for start, end in pairwise(offsets):
127 end -= LENGTH_FIELD_SIZE
128 output += bytes([0, 0, 0, 1])
129 output += data[start:end]
131 obj = cls(first_fragment=True)
132 else:
133 raise ValueError(f"NAL unit type {nal_type} is not supported")
135 return obj, output
138class H264Decoder(Decoder):
139 """H.264 receive codec. The webcam uplink is decoded by pixelflux off the
140 GIL, so this registry entry carries no libav decode of its own."""
143class H264Encoder(Encoder):
144 def __init__(self) -> None:
145 self.__target_bitrate = DEFAULT_BITRATE
147 @staticmethod
148 def _packetize_fu_a(data: Union[bytes, memoryview]) -> list[bytes]:
149 available_size = PACKET_MAX - FU_A_HEADER_SIZE
150 payload_size = len(data) - NAL_HEADER_SIZE
151 num_packets = math.ceil(payload_size / available_size)
152 num_larger_packets = payload_size % num_packets
153 package_size = payload_size // num_packets
155 f_nri = data[0] & (0x80 | 0x60) # fni of original header
156 nal = data[0] & 0x1F
158 fu_indicator = f_nri | NAL_TYPE_FU_A
160 fu_header_end = bytes([fu_indicator, nal | 0x40])
161 fu_header_middle = bytes([fu_indicator, nal])
162 fu_header_start = bytes([fu_indicator, nal | 0x80])
163 fu_header = fu_header_start
165 packages = []
166 offset = NAL_HEADER_SIZE
167 while offset < len(data):
168 if num_larger_packets > 0:
169 num_larger_packets -= 1
170 payload = data[offset : offset + package_size + 1]
171 offset += package_size + 1
172 else:
173 payload = data[offset : offset + package_size]
174 offset += package_size
176 if offset == len(data):
177 fu_header = fu_header_end
179 # join accepts a view slice, so the fragment's payload bytes are
180 # copied exactly once, straight into the outgoing RTP payload.
181 packages.append(b"".join((fu_header, payload)))
183 fu_header = fu_header_middle
184 assert offset == len(data), "incorrect fragment data"
186 return packages
188 @staticmethod
189 def _packetize_stap_a(
190 data: Union[bytes, memoryview],
191 packages_iterator: Iterator[Union[bytes, memoryview]],
192 ) -> tuple[bytes, Optional[Union[bytes, memoryview]]]:
193 counter = 0
194 available_size = PACKET_MAX - STAP_A_HEADER_SIZE
196 stap_header = NAL_TYPE_STAP_A | (data[0] & 0xE0)
198 payload = bytearray()
199 try:
200 nalu = data # with header
201 while len(nalu) <= available_size and counter < 9:
202 stap_header |= nalu[0] & 0x80
204 nri = nalu[0] & 0x60
205 if stap_header & 0x60 < nri:
206 stap_header = stap_header & 0x9F | nri
208 available_size -= LENGTH_FIELD_SIZE + len(nalu)
209 counter += 1
210 payload += pack("!H", len(nalu))
211 payload += nalu
212 nalu = next(packages_iterator)
214 if counter == 0:
215 nalu = next(packages_iterator)
216 except StopIteration:
217 nalu = None
219 # Returned payloads must be bytes: a view here would pin the whole
220 # source frame in the sender's retransmission queue.
221 if counter <= 1:
222 return (data if isinstance(data, bytes) else bytes(data)), nalu
223 else:
224 return b"".join((bytes([stap_header]), payload)), nalu
226 @staticmethod
227 def _split_bitstream(buf: Union[bytes, memoryview]) -> Iterator[memoryview]:
228 # Zero-copy NAL walk. NAL units start with the 3-byte start code
229 # 0x000001 or the 4-byte 0x00000001 (a zero byte preceding the 3-byte
230 # form). The scan runs over a memoryview and each NAL is yielded as a
231 # view slice; the packetizers copy payload bytes exactly once, when
232 # joining them with their RTP headers. The start code cannot overlap
233 # itself, so one non-overlapping scan finds every boundary.
234 mv = buf if isinstance(buf, memoryview) else memoryview(buf)
235 # An empty buffer yields no NAL units, which re would refuse to scan anyway.
236 if len(mv) == 0:
237 return
238 starts = [match.start() + 3 for match in NAL_START_CODE.finditer(mv)]
239 last = len(starts) - 1
240 for idx, nal_start in enumerate(starts):
241 if idx == last:
242 # End of buffer terminates the final NAL.
243 yield mv[nal_start:]
244 return
245 end = starts[idx + 1] - 3
246 if mv[end - 1] == 0:
247 # 4-byte start code: its leading zero is not part of this NAL.
248 end -= 1
249 yield mv[nal_start:end]
251 @classmethod
252 def _packetize(cls, packages: Iterable[Union[bytes, memoryview]]) -> list[bytes]:
253 packetized_packages = []
255 packages_iterator = iter(packages)
256 package = next(packages_iterator, None)
257 while package is not None:
258 if len(package) > PACKET_MAX:
259 packetized_packages.extend(cls._packetize_fu_a(package))
260 package = next(packages_iterator, None)
261 else:
262 packetized, package = cls._packetize_stap_a(package, packages_iterator)
263 packetized_packages.append(packetized)
265 return packetized_packages
267 def pack(self, packet: EncodedPacket) -> tuple[list[bytes], int, bool]:
268 # The encoder buffer is walked as a memoryview, so the Annex-B scan needs
269 # no full-frame copy; the views are consumed inside _packetize, which cuts
270 # them into the RTP payloads (the only bytes materialized).
271 packages = list(self._split_bitstream(memoryview(packet.data)))
272 timestamp = convert_timebase(packet.pts, packet.time_base, VIDEO_TIME_BASE)
273 keyframe = any(nal[0] & 0x1F in (5, 7) for nal in packages)
274 return self._packetize(packages), timestamp, keyframe
276 @property
277 def target_bitrate(self) -> int:
278 """
279 Target bitrate in bits per second.
280 """
281 return self.__target_bitrate
283 @target_bitrate.setter
284 def target_bitrate(self, bitrate: int) -> None:
285 bitrate = max(MIN_BITRATE, min(bitrate, MAX_BITRATE))
286 self.__target_bitrate = bitrate
289def h264_depayload(payload: bytes) -> bytes:
290 descriptor, data = H264PayloadDescriptor.parse(payload)
291 return data