codekingpro/portable-devtools
114k
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: bytes | Image.SupportsArrayInterface) -> 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 = self.fd.read(1)[0]
65 value: bytes | bytearray
66 if byte == 0b11111110 and self._previous_pixel: # QOI_OP_RGB
67 value = bytearray(self.fd.read(3)) + self._previous_pixel[3:]
68 elif byte == 0b11111111: # QOI_OP_RGBA
69 value = self.fd.read(4)
70 else:
71 op = byte >> 6
72 if op == 0: # QOI_OP_INDEX
73 op_index = byte & 0b00111111
74 value = self._previously_seen_pixels.get(
75 op_index, bytearray((0, 0, 0, 0))
76 )
77 elif op == 1 and self._previous_pixel: # QOI_OP_DIFF
78 value = bytearray(
79 (
80 (self._previous_pixel[0] + ((byte & 0b00110000) >> 4) - 2)
81 % 256,
82 (self._previous_pixel[1] + ((byte & 0b00001100) >> 2) - 2)
83 % 256,
84 (self._previous_pixel[2] + (byte & 0b00000011) - 2) % 256,
85 self._previous_pixel[3],
86 )
87 )
88 elif op == 2 and self._previous_pixel: # QOI_OP_LUMA
89 second_byte = self.fd.read(1)[0]
90 diff_green = (byte & 0b00111111) - 32
91 diff_red = ((second_byte & 0b11110000) >> 4) - 8
92 diff_blue = (second_byte & 0b00001111) - 8
93
94 value = bytearray(
95 tuple(
96 (self._previous_pixel[i] + diff_green + diff) % 256
97 for i, diff in enumerate((diff_red, 0, diff_blue))
98 )
99 )
100 value += self._previous_pixel[3:]
101 elif op == 3 and self._previous_pixel: # QOI_OP_RUN
102 run_length = (byte & 0b00111111) + 1
103 value = self._previous_pixel
104 if bands == 3:
105 value = value[:3]
106 data += value * run_length
107 continue
108 self._add_to_previous_pixels(value)
109
110 if bands == 3:
111 value = value[:3]
112 data += value
113 self.set_as_raw(data)
114 return -1, 0
115
116
117def _save(im: Image.Image, fp: IO[bytes], filename: str | bytes) -> None:
118 if im.mode == "RGB":
119 channels = 3
120 elif im.mode == "RGBA":
121 channels = 4
122 else:
123 msg = "Unsupported QOI image mode"
124 raise ValueError(msg)
125
126 colorspace = 0 if im.encoderinfo.get("colorspace") == "sRGB" else 1
127
128 fp.write(b"qoif")
129 fp.write(o32(im.size[0]))
130 fp.write(o32(im.size[1]))
131 fp.write(o8(channels))
132 fp.write(o8(colorspace))
133
134 ImageFile._save(im, fp, [ImageFile._Tile("qoi", (0, 0) + im.size)])
135
136
137class QoiEncoder(ImageFile.PyEncoder):
138 _pushes_fd = True
139 _previous_pixel: tuple[int, int, int, int] | None = None
140 _previously_seen_pixels: dict[int, tuple[int, int, int, int]] = {}
141 _run = 0
142
143 def _write_run(self) -> bytes:
144 data = o8(0b11000000 | (self._run - 1)) # QOI_OP_RUN
145 self._run = 0
146 return data
147
148 def _delta(self, left: int, right: int) -> int:
149 result = (left - right) & 255
150 if result >= 128:
151 result -= 256
152 return result
153
154 def encode(self, bufsize: int) -> tuple[int, int, bytes]:
155 assert self.im is not None
156
157 self._previously_seen_pixels = {0: (0, 0, 0, 0)}
158 self._previous_pixel = (0, 0, 0, 255)
159
160 data = bytearray()
161 w, h = self.im.size
162 bands = Image.getmodebands(self.mode)
163
164 for y in range(h):
165 for x in range(w):
166 pixel = self.im.getpixel((x, y))
167 if bands == 3:
168 pixel = (*pixel, 255)
169
170 if pixel == self._previous_pixel:
171 self._run += 1
172 if self._run == 62:
173 data += self._write_run()
174 else:
175 if self._run:
176 data += self._write_run()
177
178 r, g, b, a = pixel
179 hash_value = (r * 3 + g * 5 + b * 7 + a * 11) % 64
180 if self._previously_seen_pixels.get(hash_value) == pixel:
181 data += o8(hash_value) # QOI_OP_INDEX
182 elif self._previous_pixel:
183 self._previously_seen_pixels[hash_value] = pixel
184
185 prev_r, prev_g, prev_b, prev_a = self._previous_pixel
186 if prev_a == a:
187 delta_r = self._delta(r, prev_r)
188 delta_g = self._delta(g, prev_g)
189 delta_b = self._delta(b, prev_b)
190
191 if (
192 -2 <= delta_r < 2
193 and -2 <= delta_g < 2
194 and -2 <= delta_b < 2
195 ):
196 data += o8(
197 0b01000000
198 | (delta_r + 2) << 4
199 | (delta_g + 2) << 2
200 | (delta_b + 2)
201 ) # QOI_OP_DIFF
202 else:
203 delta_gr = self._delta(delta_r, delta_g)
204 delta_gb = self._delta(delta_b, delta_g)
205 if (
206 -8 <= delta_gr < 8
207 and -32 <= delta_g < 32
208 and -8 <= delta_gb < 8
209 ):
210 data += o8(
211 0b10000000 | (delta_g + 32)
212 ) # QOI_OP_LUMA
213 data += o8((delta_gr + 8) << 4 | (delta_gb + 8))
214 else:
215 data += o8(0b11111110) # QOI_OP_RGB
216 data += bytes(pixel[:3])
217 else:
218 data += o8(0b11111111) # QOI_OP_RGBA
219 data += bytes(pixel)
220
221 self._previous_pixel = pixel
222
223 if self._run:
224 data += self._write_run()
225 data += bytes((0, 0, 0, 0, 0, 0, 0, 1)) # padding
226
227 return len(data), 0, data
228
229
230Image.register_open(QoiImageFile.format, QoiImageFile, _accept)
231Image.register_decoder("qoi", QoiDecoder)
232Image.register_extension(QoiImageFile.format, ".qoi")
233
234Image.register_save(QoiImageFile.format, _save)
235Image.register_encoder("qoi", QoiEncoder)
236