ImagePalette.py 8.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285
  1. #
  2. # The Python Imaging Library.
  3. # $Id$
  4. #
  5. # image palette object
  6. #
  7. # History:
  8. # 1996-03-11 fl Rewritten.
  9. # 1997-01-03 fl Up and running.
  10. # 1997-08-23 fl Added load hack
  11. # 2001-04-16 fl Fixed randint shadow bug in random()
  12. #
  13. # Copyright (c) 1997-2001 by Secret Labs AB
  14. # Copyright (c) 1996-1997 by Fredrik Lundh
  15. #
  16. # See the README file for information on usage and redistribution.
  17. #
  18. from __future__ import annotations
  19. import array
  20. from collections.abc import Sequence
  21. from typing import IO, TYPE_CHECKING
  22. from . import GimpGradientFile, GimpPaletteFile, ImageColor, PaletteFile
  23. if TYPE_CHECKING:
  24. from . import Image
  25. class ImagePalette:
  26. """
  27. Color palette for palette mapped images
  28. :param mode: The mode to use for the palette. See:
  29. :ref:`concept-modes`. Defaults to "RGB"
  30. :param palette: An optional palette. If given, it must be a bytearray,
  31. an array or a list of ints between 0-255. The list must consist of
  32. all channels for one color followed by the next color (e.g. RGBRGBRGB).
  33. Defaults to an empty palette.
  34. """
  35. def __init__(
  36. self,
  37. mode: str = "RGB",
  38. palette: Sequence[int] | bytes | bytearray | None = None,
  39. ) -> None:
  40. self.mode = mode
  41. self.rawmode: str | None = None # if set, palette contains raw data
  42. self.palette = palette or bytearray()
  43. self.dirty: int | None = None
  44. @property
  45. def palette(self) -> Sequence[int] | bytes | bytearray:
  46. return self._palette
  47. @palette.setter
  48. def palette(self, palette: Sequence[int] | bytes | bytearray) -> None:
  49. self._colors: dict[tuple[int, ...], int] | None = None
  50. self._palette = palette
  51. @property
  52. def colors(self) -> dict[tuple[int, ...], int]:
  53. if self._colors is None:
  54. mode_len = len(self.mode)
  55. self._colors = {}
  56. for i in range(0, len(self.palette), mode_len):
  57. color = tuple(self.palette[i : i + mode_len])
  58. if color in self._colors:
  59. continue
  60. self._colors[color] = i // mode_len
  61. return self._colors
  62. @colors.setter
  63. def colors(self, colors: dict[tuple[int, ...], int]) -> None:
  64. self._colors = colors
  65. def copy(self) -> ImagePalette:
  66. new = ImagePalette()
  67. new.mode = self.mode
  68. new.rawmode = self.rawmode
  69. if self.palette is not None:
  70. new.palette = self.palette[:]
  71. new.dirty = self.dirty
  72. return new
  73. def getdata(self) -> tuple[str, Sequence[int] | bytes | bytearray]:
  74. """
  75. Get palette contents in format suitable for the low-level
  76. ``im.putpalette`` primitive.
  77. .. warning:: This method is experimental.
  78. """
  79. if self.rawmode:
  80. return self.rawmode, self.palette
  81. return self.mode, self.tobytes()
  82. def tobytes(self) -> bytes:
  83. """Convert palette to bytes.
  84. .. warning:: This method is experimental.
  85. """
  86. if self.rawmode:
  87. msg = "palette contains raw palette data"
  88. raise ValueError(msg)
  89. if isinstance(self.palette, bytes):
  90. return self.palette
  91. arr = array.array("B", self.palette)
  92. return arr.tobytes()
  93. # Declare tostring as an alias for tobytes
  94. tostring = tobytes
  95. def _new_color_index(
  96. self, image: Image.Image | None = None, e: Exception | None = None
  97. ) -> int:
  98. if not isinstance(self.palette, bytearray):
  99. self._palette = bytearray(self.palette)
  100. index = len(self.palette) // 3
  101. special_colors: tuple[int | tuple[int, ...] | None, ...] = ()
  102. if image:
  103. special_colors = (
  104. image.info.get("background"),
  105. image.info.get("transparency"),
  106. )
  107. while index in special_colors:
  108. index += 1
  109. if index >= 256:
  110. if image:
  111. # Search for an unused index
  112. for i, count in reversed(list(enumerate(image.histogram()))):
  113. if count == 0 and i not in special_colors:
  114. index = i
  115. break
  116. if index >= 256:
  117. msg = "cannot allocate more than 256 colors"
  118. raise ValueError(msg) from e
  119. return index
  120. def getcolor(
  121. self,
  122. color: tuple[int, ...],
  123. image: Image.Image | None = None,
  124. ) -> int:
  125. """Given an rgb tuple, allocate palette entry.
  126. .. warning:: This method is experimental.
  127. """
  128. if self.rawmode:
  129. msg = "palette contains raw palette data"
  130. raise ValueError(msg)
  131. if isinstance(color, tuple):
  132. if self.mode == "RGB":
  133. if len(color) == 4:
  134. if color[3] != 255:
  135. msg = "cannot add non-opaque RGBA color to RGB palette"
  136. raise ValueError(msg)
  137. color = color[:3]
  138. elif self.mode == "RGBA":
  139. if len(color) == 3:
  140. color += (255,)
  141. try:
  142. return self.colors[color]
  143. except KeyError as e:
  144. # allocate new color slot
  145. index = self._new_color_index(image, e)
  146. assert isinstance(self._palette, bytearray)
  147. self.colors[color] = index
  148. if index * 3 < len(self.palette):
  149. self._palette = (
  150. self._palette[: index * 3]
  151. + bytes(color)
  152. + self._palette[index * 3 + 3 :]
  153. )
  154. else:
  155. self._palette += bytes(color)
  156. self.dirty = 1
  157. return index
  158. else:
  159. msg = f"unknown color specifier: {repr(color)}" # type: ignore[unreachable]
  160. raise ValueError(msg)
  161. def save(self, fp: str | IO[str]) -> None:
  162. """Save palette to text file.
  163. .. warning:: This method is experimental.
  164. """
  165. if self.rawmode:
  166. msg = "palette contains raw palette data"
  167. raise ValueError(msg)
  168. if isinstance(fp, str):
  169. fp = open(fp, "w")
  170. fp.write("# Palette\n")
  171. fp.write(f"# Mode: {self.mode}\n")
  172. for i in range(256):
  173. fp.write(f"{i}")
  174. for j in range(i * len(self.mode), (i + 1) * len(self.mode)):
  175. try:
  176. fp.write(f" {self.palette[j]}")
  177. except IndexError:
  178. fp.write(" 0")
  179. fp.write("\n")
  180. fp.close()
  181. # --------------------------------------------------------------------
  182. # Internal
  183. def raw(rawmode: str, data: Sequence[int] | bytes | bytearray) -> ImagePalette:
  184. palette = ImagePalette()
  185. palette.rawmode = rawmode
  186. palette.palette = data
  187. palette.dirty = 1
  188. return palette
  189. # --------------------------------------------------------------------
  190. # Factories
  191. def make_linear_lut(black: int, white: float) -> list[int]:
  192. if black == 0:
  193. return [int(white * i // 255) for i in range(256)]
  194. msg = "unavailable when black is non-zero"
  195. raise NotImplementedError(msg) # FIXME
  196. def make_gamma_lut(exp: float) -> list[int]:
  197. return [int(((i / 255.0) ** exp) * 255.0 + 0.5) for i in range(256)]
  198. def negative(mode: str = "RGB") -> ImagePalette:
  199. palette = list(range(256 * len(mode)))
  200. palette.reverse()
  201. return ImagePalette(mode, [i // len(mode) for i in palette])
  202. def random(mode: str = "RGB") -> ImagePalette:
  203. from random import randint
  204. palette = [randint(0, 255) for _ in range(256 * len(mode))]
  205. return ImagePalette(mode, palette)
  206. def sepia(white: str = "#fff0c0") -> ImagePalette:
  207. bands = [make_linear_lut(0, band) for band in ImageColor.getrgb(white)]
  208. return ImagePalette("RGB", [bands[i % 3][i // 3] for i in range(256 * 3)])
  209. def wedge(mode: str = "RGB") -> ImagePalette:
  210. palette = list(range(256 * len(mode)))
  211. return ImagePalette(mode, [i // len(mode) for i in palette])
  212. def load(filename: str) -> tuple[bytes, str]:
  213. # FIXME: supports GIMP gradients only
  214. with open(filename, "rb") as fp:
  215. paletteHandlers: list[
  216. type[
  217. GimpPaletteFile.GimpPaletteFile
  218. | GimpGradientFile.GimpGradientFile
  219. | PaletteFile.PaletteFile
  220. ]
  221. ] = [
  222. GimpPaletteFile.GimpPaletteFile,
  223. GimpGradientFile.GimpGradientFile,
  224. PaletteFile.PaletteFile,
  225. ]
  226. for paletteHandler in paletteHandlers:
  227. try:
  228. fp.seek(0)
  229. lut = paletteHandler(fp).getpalette()
  230. if lut:
  231. break
  232. except (SyntaxError, ValueError):
  233. pass
  234. else:
  235. msg = "cannot load palette"
  236. raise OSError(msg)
  237. return lut # data, rawmode