Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/PIL/QoiImagePlugin.py: 53%

1# 

2# The Python Imaging Library. 

3# 

4# QOI support for PIL 

5# 

6# See the README file for information on usage and redistribution. 

7# 

8from __future__ import annotations 

9 

10import os 

11from typing import IO 

12 

13from . import Image, ImageFile 

14from ._binary import i32be as i32 

15from ._binary import o8 

16from ._binary import o32be as o32 

17 

18 

19def _accept(prefix: bytes) -> bool: 

20 return prefix.startswith(b"qoif") 

21 

22 

23class QoiImageFile(ImageFile.ImageFile): 

24 format = "QOI" 

25 format_description = "Quite OK Image" 

26 

27 def _open(self) -> None: 

28 assert self.fp is not None 

29 if not _accept(self.fp.read(4)): 

30 msg = "not a QOI file" 

31 raise SyntaxError(msg) 

32 

33 self._size = i32(self.fp.read(4)), i32(self.fp.read(4)) 

34 

35 channels = self.fp.read(1)[0] 

36 self._mode = "RGB" if channels == 3 else "RGBA" 

37 

38 self.fp.seek(1, os.SEEK_CUR) # colorspace 

39 self.tile = [ImageFile._Tile("qoi", (0, 0) + self._size, self.fp.tell())] 

40 

41 

42class QoiDecoder(ImageFile.PyDecoder): 

43 _pulls_fd = True 

44 _previous_pixel: bytes | bytearray | None = None 

45 _previously_seen_pixels: dict[int, bytes | bytearray] = {} 

46 

47 def _add_to_previous_pixels(self, value: bytes | bytearray) -> None: 

48 self._previous_pixel = value 

49 

50 r, g, b, a = value 

51 hash_value = (r * 3 + g * 5 + b * 7 + a * 11) % 64 

52 self._previously_seen_pixels[hash_value] = value 

53 

54 def decode(self, buffer: Image.DecoderInput) -> tuple[int, int]: 

55 assert self.fd is not None 

56 

57 self._previously_seen_pixels = {} 

58 self._previous_pixel = bytearray((0, 0, 0, 255)) 

59 

60 data = bytearray() 

61 bands = Image.getmodebands(self.mode) 

62 dest_length = self.state.xsize * self.state.ysize * bands 

63 while len(data) < dest_length: 

64 byte_data = self.fd.read(1) 

65 if not byte_data: 

66 break 

67 byte = byte_data[0] 

68 value: bytes | bytearray 

69 if byte == 0b11111110 and self._previous_pixel: # QOI_OP_RGB 

70 rgb_data = self.fd.read(3) 

71 if len(rgb_data) < 3: 

72 break 

73 value = bytearray(rgb_data) + self._previous_pixel[3:] 

74 elif byte == 0b11111111: # QOI_OP_RGBA 

75 value = self.fd.read(4) 

76 if len(value) < 4: 

77 break 

78 else: 

79 op = byte >> 6 

80 if op == 0: # QOI_OP_INDEX 

81 op_index = byte & 0b00111111 

82 value = self._previously_seen_pixels.get( 

83 op_index, bytearray((0, 0, 0, 0)) 

84 ) 

85 elif op == 1 and self._previous_pixel: # QOI_OP_DIFF 

86 value = bytearray( 

87 ( 

88 (self._previous_pixel[0] + ((byte & 0b00110000) >> 4) - 2) 

89 % 256, 

90 (self._previous_pixel[1] + ((byte & 0b00001100) >> 2) - 2) 

91 % 256, 

92 (self._previous_pixel[2] + (byte & 0b00000011) - 2) % 256, 

93 self._previous_pixel[3], 

94 ) 

95 ) 

96 elif op == 2 and self._previous_pixel: # QOI_OP_LUMA 

97 second_byte_data = self.fd.read(1) 

98 if not second_byte_data: 

99 break 

100 second_byte = second_byte_data[0] 

101 diff_green = (byte & 0b00111111) - 32 

102 diff_red = ((second_byte & 0b11110000) >> 4) - 8 

103 diff_blue = (second_byte & 0b00001111) - 8 

104 

105 value = bytearray( 

106 tuple( 

107 (self._previous_pixel[i] + diff_green + diff) % 256 

108 for i, diff in enumerate((diff_red, 0, diff_blue)) 

109 ) 

110 ) 

111 value += self._previous_pixel[3:] 

112 elif op == 3 and self._previous_pixel: # QOI_OP_RUN 

113 run_length = (byte & 0b00111111) + 1 

114 value = self._previous_pixel 

115 if bands == 3: 

116 value = value[:3] 

117 data += value * run_length 

118 continue 

119 self._add_to_previous_pixels(value) 

120 

121 if bands == 3: 

122 value = value[:3] 

123 data += value 

124 self.set_as_raw(data) 

125 return -1, 0 

126 

127 

128def _save(im: Image.Image, fp: IO[bytes], filename: str | bytes) -> None: 

129 if im.mode == "RGB": 

130 channels = 3 

131 elif im.mode == "RGBA": 

132 channels = 4 

133 else: 

134 msg = "Unsupported QOI image mode" 

135 raise ValueError(msg) 

136 

137 colorspace = 0 if im.encoderinfo.get("colorspace") == "sRGB" else 1 

138 

139 fp.write(b"qoif") 

140 fp.write(o32(im.size[0])) 

141 fp.write(o32(im.size[1])) 

142 fp.write(o8(channels)) 

143 fp.write(o8(colorspace)) 

144 

145 ImageFile._save(im, fp, [ImageFile._Tile("qoi", (0, 0) + im.size)]) 

146 

147 

148class QoiEncoder(ImageFile.PyEncoder): 

149 _pushes_fd = True 

150 _previous_pixel: tuple[int, int, int, int] | None = None 

151 _previously_seen_pixels: dict[int, tuple[int, int, int, int]] = {} 

152 _run = 0 

153 

154 def _write_run(self) -> bytes: 

155 data = o8(0b11000000 | (self._run - 1)) # QOI_OP_RUN 

156 self._run = 0 

157 return data 

158 

159 def _delta(self, left: int, right: int) -> int: 

160 result = (left - right) & 255 

161 if result >= 128: 

162 result -= 256 

163 return result 

164 

165 def encode(self, bufsize: int) -> tuple[int, int, bytes]: 

166 assert self.im is not None 

167 

168 self._previously_seen_pixels = {0: (0, 0, 0, 0)} 

169 self._previous_pixel = (0, 0, 0, 255) 

170 

171 data = bytearray() 

172 w, h = self.im.size 

173 bands = Image.getmodebands(self.mode) 

174 

175 for y in range(h): 

176 for x in range(w): 

177 pixel = self.im.getpixel((x, y)) 

178 if bands == 3: 

179 pixel = (*pixel, 255) 

180 

181 if pixel == self._previous_pixel: 

182 self._run += 1 

183 if self._run == 62: 

184 data += self._write_run() 

185 else: 

186 if self._run: 

187 data += self._write_run() 

188 

189 r, g, b, a = pixel 

190 hash_value = (r * 3 + g * 5 + b * 7 + a * 11) % 64 

191 if self._previously_seen_pixels.get(hash_value) == pixel: 

192 data += o8(hash_value) # QOI_OP_INDEX 

193 elif self._previous_pixel: 

194 self._previously_seen_pixels[hash_value] = pixel 

195 

196 prev_r, prev_g, prev_b, prev_a = self._previous_pixel 

197 if prev_a == a: 

198 delta_r = self._delta(r, prev_r) 

199 delta_g = self._delta(g, prev_g) 

200 delta_b = self._delta(b, prev_b) 

201 

202 if ( 

203 -2 <= delta_r < 2 

204 and -2 <= delta_g < 2 

205 and -2 <= delta_b < 2 

206 ): 

207 data += o8( 

208 0b01000000 

209 | (delta_r + 2) << 4 

210 | (delta_g + 2) << 2 

211 | (delta_b + 2) 

212 ) # QOI_OP_DIFF 

213 else: 

214 delta_gr = self._delta(delta_r, delta_g) 

215 delta_gb = self._delta(delta_b, delta_g) 

216 if ( 

217 -8 <= delta_gr < 8 

218 and -32 <= delta_g < 32 

219 and -8 <= delta_gb < 8 

220 ): 

221 data += o8( 

222 0b10000000 | (delta_g + 32) 

223 ) # QOI_OP_LUMA 

224 data += o8((delta_gr + 8) << 4 | (delta_gb + 8)) 

225 else: 

226 data += o8(0b11111110) # QOI_OP_RGB 

227 data += bytes(pixel[:3]) 

228 else: 

229 data += o8(0b11111111) # QOI_OP_RGBA 

230 data += bytes(pixel) 

231 

232 self._previous_pixel = pixel 

233 

234 if self._run: 

235 data += self._write_run() 

236 data += bytes((0, 0, 0, 0, 0, 0, 0, 1)) # padding 

237 

238 return len(data), 0, bytes(data) 

239 

240 

241Image.register_open(QoiImageFile.format, QoiImageFile, _accept) 

242Image.register_decoder("qoi", QoiDecoder) 

243Image.register_extension(QoiImageFile.format, ".qoi") 

244 

245Image.register_save(QoiImageFile.format, _save) 

246Image.register_encoder("qoi", QoiEncoder)