ImagePalette.py 7.4 KB

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