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

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 9628 packetization of VP9 frames in the non-flexible mode a single-layer 

7stream takes: a 15-bit picture id on every packet, the inter-picture flag off 

8on key frames, which also open with the scalability structure (one spatial 

9layer at the frame's size, one temporal layer, each picture predicted from 

10the one before), and the frame's first and last packets marked.""" 

11 

12import random 

13from struct import pack 

14from typing import Optional, Union 

15 

16from ..mediastreams import VIDEO_TIME_BASE, convert_timebase 

17from .base import Decoder, Encoder, EncodedPacket 

18from .h264 import PACKET_MAX 

19 

20Buffer = Union[bytes, memoryview] 

21 

22FLAG_I = 0x80 # picture id present 

23FLAG_P = 0x40 # inter-picture predicted 

24FLAG_L = 0x20 # layer indices present 

25FLAG_F = 0x10 # flexible mode 

26FLAG_B = 0x08 # first packet of a frame 

27FLAG_E = 0x04 # last packet of a frame 

28FLAG_V = 0x02 # scalability structure present 

29FLAG_Z = 0x01 # not a reference for upper spatial layers 

30 

31VP9_SYNC_CODE = (0x49, 0x83, 0x42) 

32 

33 

34class _BitReader: 

35 def __init__(self, data: Buffer) -> None: 

36 self.data = data 

37 self.pos = 0 

38 

39 def u(self, bits: int) -> int: 

40 value = 0 

41 for _ in range(bits): 

42 byte = self.data[self.pos >> 3] 

43 value = (value << 1) | ((byte >> (7 - (self.pos & 7))) & 1) 

44 self.pos += 1 

45 return value 

46 

47 

48def vp9_profile(frame: Buffer) -> int: 

49 b = frame[0] 

50 return ((b >> 5) & 1) | (((b >> 4) & 1) << 1) 

51 

52 

53def vp9_is_key(frame: Buffer) -> bool: 

54 """Read from the uncompressed header: not a shown existing frame, and 

55 `frame_type` KEY_FRAME. Profile 3 carries a reserved bit ahead of both.""" 

56 if not frame or frame[0] >> 6 != 0b10: 

57 return False 

58 b = frame[0] 

59 if vp9_profile(frame) == 3: 

60 return b & 0x04 == 0 and b & 0x02 == 0 

61 return b & 0x08 == 0 and b & 0x04 == 0 

62 

63 

64def vp9_key_frame_size(frame: Buffer) -> Optional[tuple[int, int]]: 

65 """The frame size a key frame's uncompressed header states, or None for 

66 anything that is not a parseable key frame.""" 

67 if len(frame) < 10 or not vp9_is_key(frame): 

68 return None 

69 try: 

70 r = _BitReader(frame) 

71 r.u(2) 

72 profile = r.u(1) | (r.u(1) << 1) 

73 if profile == 3: 

74 r.u(1) 

75 r.u(1) # show_existing_frame 

76 r.u(1) # frame_type 

77 r.u(1) # show_frame 

78 r.u(1) # error_resilient_mode 

79 if (r.u(8), r.u(8), r.u(8)) != VP9_SYNC_CODE: 

80 return None 

81 if profile >= 2: 

82 r.u(1) # ten_or_twelve_bit 

83 color_space = r.u(3) 

84 if color_space != 7: 

85 r.u(1) # color_range 

86 if profile in (1, 3): 

87 r.u(3) # subsampling_x, subsampling_y, reserved 

88 elif profile in (1, 3): 

89 r.u(1) 

90 width = r.u(16) + 1 

91 height = r.u(16) + 1 

92 return width, height 

93 except IndexError: 

94 return None 

95 

96 

97class Vp9Decoder(Decoder): 

98 pass 

99 

100 

101class Vp9Encoder(Encoder): 

102 def __init__(self) -> None: 

103 self.picture_id = random.randint(0, (1 << 15) - 1) 

104 

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

106 frame = memoryview(packet.data) 

107 keyframe = vp9_is_key(frame) 

108 payloads = self._packetize(frame, self.picture_id, keyframe) 

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

110 self.picture_id = (self.picture_id + 1) % (1 << 15) 

111 return payloads, timestamp, keyframe 

112 

113 @staticmethod 

114 def _scalability_structure(frame: Buffer) -> bytes: 

115 size = vp9_key_frame_size(frame) 

116 # One spatial layer, its resolution when known, one picture group whose 

117 # single picture (temporal layer 0, an up-switch point) references the 

118 # picture before it. 

119 ss = bytearray([0x08 | (0x10 if size else 0)]) 

120 if size: 

121 ss += pack("!HH", size[0], size[1]) 

122 ss += bytes([1, 0x14, 1]) 

123 return bytes(ss) 

124 

125 @classmethod 

126 def _packetize(cls, frame: Buffer, picture_id: int, keyframe: bool) -> list[bytes]: 

127 pid = pack("!H", 0x8000 | picture_id) 

128 ss = cls._scalability_structure(frame) if keyframe else b"" 

129 payloads = [] 

130 length = len(frame) 

131 pos = 0 

132 first = True 

133 while pos < length or first: 

134 flags = FLAG_I | (0 if keyframe else FLAG_P) 

135 if first: 

136 flags |= FLAG_B | (FLAG_V if keyframe else 0) 

137 descriptor = bytes([flags]) + pid + (ss if first else b"") 

138 size = min(length - pos, PACKET_MAX - len(descriptor)) 

139 end = pos + size 

140 if end >= length: 

141 descriptor = bytes([flags | FLAG_E]) + descriptor[1:] 

142 payloads.append(b"".join((descriptor, frame[pos:end]))) 

143 pos = end 

144 first = False 

145 return payloads 

146 

147 

148def vp9_depayload(payload: bytes) -> bytes: 

149 if len(payload) < 1: 

150 raise ValueError("VP9 descriptor is too short") 

151 flags = payload[0] 

152 pos = 1 

153 try: 

154 if flags & FLAG_I: 

155 pos += 2 if payload[pos] & 0x80 else 1 

156 if flags & FLAG_L: 

157 pos += 1 

158 if not flags & FLAG_F: 

159 pos += 1 # TL0PICIDX 

160 if flags & FLAG_F and flags & FLAG_P: 

161 while True: 

162 more = payload[pos] & 1 

163 pos += 1 

164 if not more: 

165 break 

166 if flags & FLAG_V: 

167 b = payload[pos] 

168 pos += 1 

169 layers = (b >> 5) + 1 

170 if b & 0x10: 

171 pos += 4 * layers 

172 if b & 0x08: 

173 groups = payload[pos] 

174 pos += 1 

175 for _ in range(groups): 

176 pos += 1 + ((payload[pos] >> 2) & 3) 

177 except IndexError: 

178 raise ValueError("VP9 descriptor is truncated") 

179 if pos > len(payload): 

180 raise ValueError("VP9 descriptor is truncated") 

181 return payload[pos:]