ImagePalette.py 7.6 KB

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