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#
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"""AV1 RTP payload format (AOM, v1.0) packetization of temporal units, and the
7reassembly of a received frame's packets back into one.
9Packing: the temporal delimiter, tile lists, and padding go; every other OBU
10becomes an element without its size field. Elements fill packets behind a
11one-byte aggregation header (Z: the first element continues an OBU from the
12previous packet, Y: the last continues in the next, W: the element count when
13three or fewer, N: the first packet of a coded video sequence), a LEB128 length
14ahead of every element but a counted packet's last, and an OBU larger than the
15room a packet has left is split across as many packets as it needs.
17Assembly runs over a frame's packets in sequence order once the jitter buffer
18has them all: fragments are joined across packets by the Y and Z flags, each
19OBU gets the size field the elements travel without, and a temporal delimiter
20opens the unit, which is then the low-overhead bitstream the virtual camera's
21decoder reads. A fragment whose start was lost with a packet is dropped, as
22are the rest of that OBU's fragments."""
24from typing import Union
26from ..mediastreams import VIDEO_TIME_BASE, convert_timebase
27from .base import Decoder, Encoder, EncodedPacket
28from .h264 import PACKET_MAX
30Buffer = Union[bytes, memoryview]
32OBU_SEQUENCE_HEADER = 1
33OBU_TEMPORAL_DELIMITER = 2
34OBU_FRAME_HEADER = 3
35OBU_FRAME = 6
36OBU_TILE_LIST = 8
37OBU_PADDING = 15
38DROPPED_OBUS = (OBU_TEMPORAL_DELIMITER, OBU_TILE_LIST, OBU_PADDING)
40AGGREGATION_HEADER_SIZE = 1
41# The smallest fragment worth starting in a packet's remaining room.
42MIN_FRAGMENT = 16
45def leb128(value: int) -> bytes:
46 out = bytearray()
47 while True:
48 byte = value & 0x7F
49 value >>= 7
50 if value:
51 out.append(byte | 0x80)
52 else:
53 out.append(byte)
54 return bytes(out)
57def read_leb128(data: Buffer, pos: int) -> tuple[int, int]:
58 value = 0
59 shift = 0
60 while True:
61 if pos >= len(data) or shift > 56:
62 raise ValueError("LEB128 is truncated")
63 byte = data[pos]
64 pos += 1
65 value |= (byte & 0x7F) << shift
66 if not byte & 0x80:
67 return value, pos
68 shift += 7
71def av1_obus(tu: Buffer) -> list[tuple[int, bytes, Buffer]]:
72 """The OBUs of a temporal unit as `(type, header, payload)`, the header
73 without its size flag so the element carries no size field."""
74 obus = []
75 pos = 0
76 while pos < len(tu):
77 header = tu[pos]
78 obu_type = (header >> 3) & 0x0F
79 has_extension = bool(header & 0x04)
80 has_size = bool(header & 0x02)
81 start = pos + 1 + has_extension
82 if has_size:
83 size, start = read_leb128(tu, start)
84 else:
85 size = len(tu) - start
86 end = start + size
87 if end > len(tu):
88 raise ValueError("OBU is truncated")
89 header_bytes = bytes([header & ~0x02]) + (bytes(tu[pos + 1 : pos + 2]) if has_extension else b"")
90 obus.append((obu_type, header_bytes, tu[start:end]))
91 pos = end
92 return obus
95def av1_is_key(obus: list[tuple[int, bytes, Buffer]]) -> bool:
96 """Whether the unit opens with a key frame: its first frame header or frame
97 OBU's `frame_type`, read past the temporal delimiter and sequence header.
98 A shown existing frame, or a unit without a frame, is a delta."""
99 for obu_type, _, payload in obus:
100 if obu_type in (OBU_FRAME_HEADER, OBU_FRAME):
101 if not payload:
102 return False
103 first = payload[0]
104 return not first & 0x80 and (first >> 5) & 0x03 == 0
105 return False
108class Av1Decoder(Decoder):
109 pass
112def _obu_elements(payload: bytes) -> tuple[bool, bool, list[bytes]]:
113 """A packet's Z and Y flags and its OBU elements, sizes read off."""
114 header = payload[0]
115 w = (header >> 4) & 3
116 pos = 1
117 elements: list[bytes] = []
118 while pos < len(payload):
119 if w == 0 or len(elements) < w - 1:
120 size, pos = read_leb128(payload, pos)
121 else:
122 size = len(payload) - pos
123 if pos + size > len(payload):
124 raise ValueError("OBU element is truncated")
125 elements.append(payload[pos : pos + size])
126 pos += size
127 return bool(header & 0x80), bool(header & 0x40), elements
130def _sized_obu(obu: bytes) -> bytes:
131 """An OBU element with the size field the bitstream format wants."""
132 header = obu[0]
133 header_size = 2 if header & 0x04 else 1
134 if header & 0x02 or len(obu) < header_size:
135 return obu
136 return bytes([header | 0x02]) + obu[1:header_size] + leb128(len(obu) - header_size) + obu[header_size:]
139def av1_assemble(payloads: list[bytes]) -> bytes:
140 """The temporal unit a frame's packets carry, in the low-overhead bitstream
141 format: a temporal delimiter, then every OBU element made whole across
142 packets and given its size field."""
143 out = bytearray(b"\x12\x00")
144 pending = bytearray()
145 for payload in payloads:
146 if not payload:
147 continue
148 try:
149 continues, continued, elements = _obu_elements(payload)
150 except ValueError:
151 pending = bytearray()
152 continue
153 for index, element in enumerate(elements):
154 first = index == 0
155 last = index == len(elements) - 1
156 if first and continues:
157 if not pending:
158 # The OBU's start went with a lost packet; so does its rest.
159 continue
160 pending += element
161 else:
162 pending = bytearray(element)
163 if last and continued:
164 continue
165 if pending:
166 out += _sized_obu(bytes(pending))
167 pending = bytearray()
168 return bytes(out)
171class Av1Encoder(Encoder):
172 def pack(self, packet: EncodedPacket) -> tuple[list[bytes], int, bool]:
173 obus = av1_obus(memoryview(packet.data))
174 keyframe = av1_is_key(obus)
175 elements = [b"".join((header, payload)) for obu_type, header, payload in obus
176 if obu_type not in DROPPED_OBUS]
177 timestamp = convert_timebase(packet.pts, packet.time_base, VIDEO_TIME_BASE)
178 return self._packetize(elements, keyframe), timestamp, keyframe
180 @staticmethod
181 def _packetize(elements: list[bytes], keyframe: bool) -> list[bytes]:
182 packets: list[bytes] = []
183 current: list[bytes] = []
184 used = AGGREGATION_HEADER_SIZE
185 continues_previous = False
187 def flush(continues_next: bool) -> None:
188 nonlocal current, used, continues_previous
189 count = len(current)
190 w = count if count <= 3 else 0
191 header = (
192 (0x80 if continues_previous else 0)
193 | (0x40 if continues_next else 0)
194 | (w << 4)
195 | (0x08 if keyframe and not packets else 0)
196 )
197 body = [bytes([header])]
198 for index, element in enumerate(current):
199 if w == 0 or index < count - 1:
200 body.append(leb128(len(element)))
201 body.append(element)
202 packets.append(b"".join(body))
203 current = []
204 used = AGGREGATION_HEADER_SIZE
205 continues_previous = continues_next
207 for element in elements:
208 # Every element is costed with its length prefix; a counted packet's
209 # last element drops it, which only leaves room to spare.
210 need = len(leb128(len(element))) + len(element)
211 if used + need <= PACKET_MAX:
212 current.append(element)
213 used += need
214 continue
215 room = PACKET_MAX - used - len(leb128(PACKET_MAX))
216 if room < MIN_FRAGMENT:
217 flush(False)
218 room = PACKET_MAX - used - len(leb128(PACKET_MAX))
219 pos = 0
220 while pos < len(element):
221 take = min(len(element) - pos, room)
222 current.append(element[pos : pos + take])
223 used += len(leb128(take)) + take
224 pos += take
225 if pos < len(element):
226 flush(True)
227 room = PACKET_MAX - used - len(leb128(PACKET_MAX))
228 if current or not packets:
229 flush(False)
230 return packets