WebPImagePlugin.py 9.6 KB

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