ImagePalette.py 8.8 KB

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