ImagePalette.py 7.7 KB

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