WebPImagePlugin.py 9.8 KB

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