WebPImagePlugin.py 9.8 KB

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