ImagePalette.py 7.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255
  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. import array
  19. from . import GimpGradientFile, GimpPaletteFile, ImageColor, PaletteFile
  20. from ._deprecate import deprecate
  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, size=0):
  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. if size != 0:
  37. deprecate("The size parameter", 10, None)
  38. if size != len(self.palette):
  39. raise ValueError("wrong palette size")
  40. @property
  41. def palette(self):
  42. return self._palette
  43. @palette.setter
  44. def palette(self, palette):
  45. self._palette = palette
  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. def copy(self):
  54. new = ImagePalette()
  55. new.mode = self.mode
  56. new.rawmode = self.rawmode
  57. if self.palette is not None:
  58. new.palette = self.palette[:]
  59. new.dirty = self.dirty
  60. return new
  61. def getdata(self):
  62. """
  63. Get palette contents in format suitable for the low-level
  64. ``im.putpalette`` primitive.
  65. .. warning:: This method is experimental.
  66. """
  67. if self.rawmode:
  68. return self.rawmode, self.palette
  69. return self.mode, self.tobytes()
  70. def tobytes(self):
  71. """Convert palette to bytes.
  72. .. warning:: This method is experimental.
  73. """
  74. if self.rawmode:
  75. raise ValueError("palette contains raw palette data")
  76. if isinstance(self.palette, bytes):
  77. return self.palette
  78. arr = array.array("B", self.palette)
  79. return arr.tobytes()
  80. # Declare tostring as an alias for tobytes
  81. tostring = tobytes
  82. def getcolor(self, color, image=None):
  83. """Given an rgb tuple, allocate palette entry.
  84. .. warning:: This method is experimental.
  85. """
  86. if self.rawmode:
  87. raise ValueError("palette contains raw palette data")
  88. if isinstance(color, tuple):
  89. if self.mode == "RGB":
  90. if len(color) == 4 and color[3] == 255:
  91. color = color[:3]
  92. elif self.mode == "RGBA":
  93. if len(color) == 3:
  94. color += (255,)
  95. try:
  96. return self.colors[color]
  97. except KeyError as e:
  98. # allocate new color slot
  99. if not isinstance(self.palette, bytearray):
  100. self._palette = bytearray(self.palette)
  101. index = len(self.palette) // 3
  102. special_colors = ()
  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. raise ValueError("cannot allocate more than 256 colors") from e
  119. self.colors[color] = index
  120. if index * 3 < len(self.palette):
  121. self._palette = (
  122. self.palette[: index * 3]
  123. + bytes(color)
  124. + self.palette[index * 3 + 3 :]
  125. )
  126. else:
  127. self._palette += bytes(color)
  128. self.dirty = 1
  129. return index
  130. else:
  131. raise ValueError(f"unknown color specifier: {repr(color)}")
  132. def save(self, fp):
  133. """Save palette to text file.
  134. .. warning:: This method is experimental.
  135. """
  136. if self.rawmode:
  137. raise ValueError("palette contains raw palette data")
  138. if isinstance(fp, str):
  139. fp = open(fp, "w")
  140. fp.write("# Palette\n")
  141. fp.write(f"# Mode: {self.mode}\n")
  142. for i in range(256):
  143. fp.write(f"{i}")
  144. for j in range(i * len(self.mode), (i + 1) * len(self.mode)):
  145. try:
  146. fp.write(f" {self.palette[j]}")
  147. except IndexError:
  148. fp.write(" 0")
  149. fp.write("\n")
  150. fp.close()
  151. # --------------------------------------------------------------------
  152. # Internal
  153. def raw(rawmode, data):
  154. palette = ImagePalette()
  155. palette.rawmode = rawmode
  156. palette.palette = data
  157. palette.dirty = 1
  158. return palette
  159. # --------------------------------------------------------------------
  160. # Factories
  161. def make_linear_lut(black, white):
  162. lut = []
  163. if black == 0:
  164. for i in range(256):
  165. lut.append(white * i // 255)
  166. else:
  167. raise NotImplementedError # FIXME
  168. return lut
  169. def make_gamma_lut(exp):
  170. lut = []
  171. for i in range(256):
  172. lut.append(int(((i / 255.0) ** exp) * 255.0 + 0.5))
  173. return lut
  174. def negative(mode="RGB"):
  175. palette = list(range(256 * len(mode)))
  176. palette.reverse()
  177. return ImagePalette(mode, [i // len(mode) for i in palette])
  178. def random(mode="RGB"):
  179. from random import randint
  180. palette = []
  181. for i in range(256 * len(mode)):
  182. palette.append(randint(0, 255))
  183. return ImagePalette(mode, palette)
  184. def sepia(white="#fff0c0"):
  185. bands = [make_linear_lut(0, band) for band in ImageColor.getrgb(white)]
  186. return ImagePalette("RGB", [bands[i % 3][i // 3] for i in range(256 * 3)])
  187. def wedge(mode="RGB"):
  188. palette = list(range(256 * len(mode)))
  189. return ImagePalette(mode, [i // len(mode) for i in palette])
  190. def load(filename):
  191. # FIXME: supports GIMP gradients only
  192. with open(filename, "rb") as fp:
  193. for paletteHandler in [
  194. GimpPaletteFile.GimpPaletteFile,
  195. GimpGradientFile.GimpGradientFile,
  196. PaletteFile.PaletteFile,
  197. ]:
  198. try:
  199. fp.seek(0)
  200. lut = paletteHandler(fp).getpalette()
  201. if lut:
  202. break
  203. except (SyntaxError, ValueError):
  204. # import traceback
  205. # traceback.print_exc()
  206. pass
  207. else:
  208. raise OSError("cannot load palette")
  209. return lut # data, rawmode