You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

337 lines
10 KiB

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