1from __future__ import annotations
2
3from io import BytesIO
4from typing import IO, Any
5
6from . import Image, ImageFile
7
8try:
9 from . import _webp
10
11 SUPPORTED = True
12except ImportError:
13 SUPPORTED = False
14
15
16_VP8_MODES_BY_IDENTIFIER = {
17 b"VP8 ": "RGB",
18 b"VP8X": "RGBA",
19 b"VP8L": "RGBA", # lossless
20}
21
22
23def _accept(prefix: bytes) -> bool | str:
24 is_riff_file_format = prefix[:4] == b"RIFF"
25 is_webp_file = prefix[8:12] == b"WEBP"
26 is_valid_vp8_mode = prefix[12:16] in _VP8_MODES_BY_IDENTIFIER
27
28 if is_riff_file_format and is_webp_file and is_valid_vp8_mode:
29 if not SUPPORTED:
30 return (
31 "image file could not be identified because WEBP support not installed"
32 )
33 return True
34 return False
35
36
37class WebPImageFile(ImageFile.ImageFile):
38 format = "WEBP"
39 format_description = "WebP image"
40 __loaded = 0
41 __logical_frame = 0
42
43 def _open(self) -> None:
44 # Use the newer AnimDecoder API to parse the (possibly) animated file,
45 # and access muxed chunks like ICC/EXIF/XMP.
46 self._decoder = _webp.WebPAnimDecoder(self.fp.read())
47
48 # Get info from decoder
49 width, height, loop_count, bgcolor, frame_count, mode = self._decoder.get_info()
50 self._size = width, height
51 self.info["loop"] = loop_count
52 bg_a, bg_r, bg_g, bg_b = (
53 (bgcolor >> 24) & 0xFF,
54 (bgcolor >> 16) & 0xFF,
55 (bgcolor >> 8) & 0xFF,
56 bgcolor & 0xFF,
57 )
58 self.info["background"] = (bg_r, bg_g, bg_b, bg_a)
59 self.n_frames = frame_count
60 self.is_animated = self.n_frames > 1
61 self._mode = "RGB" if mode == "RGBX" else mode
62 self.rawmode = mode
63
64 # Attempt to read ICC / EXIF / XMP chunks from file
65 icc_profile = self._decoder.get_chunk("ICCP")
66 exif = self._decoder.get_chunk("EXIF")
67 xmp = self._decoder.get_chunk("XMP ")
68 if icc_profile:
69 self.info["icc_profile"] = icc_profile
70 if exif:
71 self.info["exif"] = exif
72 if xmp:
73 self.info["xmp"] = xmp
74
75 # Initialize seek state
76 self._reset(reset=False)
77
78 def _getexif(self) -> dict[int, Any] | None:
79 if "exif" not in self.info:
80 return None
81 return self.getexif()._get_merged_dict()
82
83 def seek(self, frame: int) -> None:
84 if not self._seek_check(frame):
85 return
86
87 # Set logical frame to requested position
88 self.__logical_frame = frame
89
90 def _reset(self, reset: bool = True) -> None:
91 if reset:
92 self._decoder.reset()
93 self.__physical_frame = 0
94 self.__loaded = -1
95 self.__timestamp = 0
96
97 def _get_next(self) -> tuple[bytes, int, int]:
98 # Get next frame
99 ret = self._decoder.get_next()
100 self.__physical_frame += 1
101
102 # Check if an error occurred
103 if ret is None:
104 self._reset() # Reset just to be safe
105 self.seek(0)
106 msg = "failed to decode next frame in WebP file"
107 raise EOFError(msg)
108
109 # Compute duration
110 data, timestamp = ret
111 duration = timestamp - self.__timestamp
112 self.__timestamp = timestamp
113
114 # libwebp gives frame end, adjust to start of frame
115 timestamp -= duration
116 return data, timestamp, duration
117
118 def _seek(self, frame: int) -> None:
119 if self.__physical_frame == frame:
120 return # Nothing to do
121 if frame < self.__physical_frame:
122 self._reset() # Rewind to beginning
123 while self.__physical_frame < frame:
124 self._get_next() # Advance to the requested frame
125
126 def load(self) -> Image.core.PixelAccess | None:
127 if self.__loaded != self.__logical_frame:
128 self._seek(self.__logical_frame)
129
130 # We need to load the image data for this frame
131 data, timestamp, duration = self._get_next()
132 self.info["timestamp"] = timestamp
133 self.info["duration"] = duration
134 self.__loaded = self.__logical_frame
135
136 # Set tile
137 if self.fp and self._exclusive_fp:
138 self.fp.close()
139 self.fp = BytesIO(data)
140 self.tile = [ImageFile._Tile("raw", (0, 0) + self.size, 0, self.rawmode)]
141
142 return super().load()
143
144 def load_seek(self, pos: int) -> None:
145 pass
146
147 def tell(self) -> int:
148 return self.__logical_frame
149
150
151def _convert_frame(im: Image.Image) -> Image.Image:
152 # Make sure image mode is supported
153 if im.mode not in ("RGBX", "RGBA", "RGB"):
154 im = im.convert("RGBA" if im.has_transparency_data else "RGB")
155 return im
156
157
158def _save_all(im: Image.Image, fp: IO[bytes], filename: str | bytes) -> None:
159 encoderinfo = im.encoderinfo.copy()
160 append_images = list(encoderinfo.get("append_images", []))
161
162 # If total frame count is 1, then save using the legacy API, which
163 # will preserve non-alpha modes
164 total = 0
165 for ims in [im] + append_images:
166 total += getattr(ims, "n_frames", 1)
167 if total == 1:
168 _save(im, fp, filename)
169 return
170
171 background: int | tuple[int, ...] = (0, 0, 0, 0)
172 if "background" in encoderinfo:
173 background = encoderinfo["background"]
174 elif "background" in im.info:
175 background = im.info["background"]
176 if isinstance(background, int):
177 # GifImagePlugin stores a global color table index in
178 # info["background"]. So it must be converted to an RGBA value
179 palette = im.getpalette()
180 if palette:
181 r, g, b = palette[background * 3 : (background + 1) * 3]
182 background = (r, g, b, 255)
183 else:
184 background = (background, background, background, 255)
185
186 duration = im.encoderinfo.get("duration", im.info.get("duration", 0))
187 loop = im.encoderinfo.get("loop", 0)
188 minimize_size = im.encoderinfo.get("minimize_size", False)
189 kmin = im.encoderinfo.get("kmin", None)
190 kmax = im.encoderinfo.get("kmax", None)
191 allow_mixed = im.encoderinfo.get("allow_mixed", False)
192 verbose = False
193 lossless = im.encoderinfo.get("lossless", False)
194 quality = im.encoderinfo.get("quality", 80)
195 alpha_quality = im.encoderinfo.get("alpha_quality", 100)
196 method = im.encoderinfo.get("method", 0)
197 icc_profile = im.encoderinfo.get("icc_profile") or ""
198 exif = im.encoderinfo.get("exif", "")
199 if isinstance(exif, Image.Exif):
200 exif = exif.tobytes()
201 xmp = im.encoderinfo.get("xmp", "")
202 if allow_mixed:
203 lossless = False
204
205 # Sensible keyframe defaults are from gif2webp.c script
206 if kmin is None:
207 kmin = 9 if lossless else 3
208 if kmax is None:
209 kmax = 17 if lossless else 5
210
211 # Validate background color
212 if (
213 not isinstance(background, (list, tuple))
214 or len(background) != 4
215 or not all(0 <= v < 256 for v in background)
216 ):
217 msg = f"Background color is not an RGBA tuple clamped to (0-255): {background}"
218 raise OSError(msg)
219
220 # Convert to packed uint
221 bg_r, bg_g, bg_b, bg_a = background
222 background = (bg_a << 24) | (bg_r << 16) | (bg_g << 8) | (bg_b << 0)
223
224 # Setup the WebP animation encoder
225 enc = _webp.WebPAnimEncoder(
226 im.size[0],
227 im.size[1],
228 background,
229 loop,
230 minimize_size,
231 kmin,
232 kmax,
233 allow_mixed,
234 verbose,
235 )
236
237 # Add each frame
238 frame_idx = 0
239 timestamp = 0
240 cur_idx = im.tell()
241 try:
242 for ims in [im] + append_images:
243 # Get # of frames in this image
244 nfr = getattr(ims, "n_frames", 1)
245
246 for idx in range(nfr):
247 ims.seek(idx)
248
249 frame = _convert_frame(ims)
250
251 # Append the frame to the animation encoder
252 enc.add(
253 frame.getim(),
254 round(timestamp),
255 lossless,
256 quality,
257 alpha_quality,
258 method,
259 )
260
261 # Update timestamp and frame index
262 if isinstance(duration, (list, tuple)):
263 timestamp += duration[frame_idx]
264 else:
265 timestamp += duration
266 frame_idx += 1
267
268 finally:
269 im.seek(cur_idx)
270
271 # Force encoder to flush frames
272 enc.add(None, round(timestamp), lossless, quality, alpha_quality, 0)
273
274 # Get the final output from the encoder
275 data = enc.assemble(icc_profile, exif, xmp)
276 if data is None:
277 msg = "cannot write file as WebP (encoder returned None)"
278 raise OSError(msg)
279
280 fp.write(data)
281
282
283def _save(im: Image.Image, fp: IO[bytes], filename: str | bytes) -> None:
284 lossless = im.encoderinfo.get("lossless", False)
285 quality = im.encoderinfo.get("quality", 80)
286 alpha_quality = im.encoderinfo.get("alpha_quality", 100)
287 icc_profile = im.encoderinfo.get("icc_profile") or ""
288 exif = im.encoderinfo.get("exif", b"")
289 if isinstance(exif, Image.Exif):
290 exif = exif.tobytes()
291 if exif.startswith(b"Exif\x00\x00"):
292 exif = exif[6:]
293 xmp = im.encoderinfo.get("xmp", "")
294 method = im.encoderinfo.get("method", 4)
295 exact = 1 if im.encoderinfo.get("exact") else 0
296
297 im = _convert_frame(im)
298
299 data = _webp.WebPEncode(
300 im.getim(),
301 lossless,
302 float(quality),
303 float(alpha_quality),
304 icc_profile,
305 method,
306 exact,
307 exif,
308 xmp,
309 )
310 if data is None:
311 msg = "cannot write file as WebP (encoder returned None)"
312 raise OSError(msg)
313
314 fp.write(data)
315
316
317Image.register_open(WebPImageFile.format, WebPImageFile, _accept)
318if SUPPORTED:
319 Image.register_save(WebPImageFile.format, _save)
320 Image.register_save_all(WebPImageFile.format, _save_all)
321 Image.register_extension(WebPImageFile.format, ".webp")
322 Image.register_mime(WebPImageFile.format, "image/webp")