images2gif.py 35 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014101510161017101810191020102110221023102410251026102710281029103010311032103310341035103610371038103910401041104210431044104510461047104810491050105110521053105410551056105710581059
  1. # -*- coding: utf-8 -*-
  2. # Copyright (C) 2012, Almar Klein, Ant1, Marius van Voorden
  3. #
  4. # This code is subject to the (new) BSD license:
  5. #
  6. # Redistribution and use in source and binary forms, with or without
  7. # modification, are permitted provided that the following conditions are met:
  8. # * Redistributions of source code must retain the above copyright
  9. # notice, this list of conditions and the following disclaimer.
  10. # * Redistributions in binary form must reproduce the above copyright
  11. # notice, this list of conditions and the following disclaimer in the
  12. # documentation and/or other materials provided with the distribution.
  13. # * Neither the name of the <organization> nor the
  14. # names of its contributors may be used to endorse or promote products
  15. # derived from this software without specific prior written permission.
  16. #
  17. # THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
  18. # AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
  19. # IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
  20. # ARE DISCLAIMED. IN NO EVENT SHALL <COPYRIGHT HOLDER> BE LIABLE FOR ANY
  21. # DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
  22. # (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
  23. # LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
  24. # ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
  25. # (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
  26. # SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
  27. """ Module images2gif
  28. Provides functionality for reading and writing animated GIF images.
  29. Use writeGif to write a series of numpy arrays or PIL images as an
  30. animated GIF. Use readGif to read an animated gif as a series of numpy
  31. arrays.
  32. Note that since July 2004, all patents on the LZW compression patent have
  33. expired. Therefore the GIF format may now be used freely.
  34. Acknowledgements
  35. ----------------
  36. Many thanks to Ant1 for:
  37. * noting the use of "palette=PIL.Image.ADAPTIVE", which significantly
  38. improves the results.
  39. * the modifications to save each image with its own palette, or optionally
  40. the global palette (if its the same).
  41. Many thanks to Marius van Voorden for porting the NeuQuant quantization
  42. algorithm of Anthony Dekker to Python (See the NeuQuant class for its
  43. license).
  44. Many thanks to Alex Robinson for implementing the concept of subrectangles,
  45. which (depening on image content) can give a very significant reduction in
  46. file size.
  47. This code is based on gifmaker (in the scripts folder of the source
  48. distribution of PIL)
  49. Usefull links
  50. -------------
  51. * http://tronche.com/computer-graphics/gif/
  52. * http://en.wikipedia.org/wiki/Graphics_Interchange_Format
  53. * http://www.w3.org/Graphics/GIF/spec-gif89a.txt
  54. """
  55. # todo: This module should be part of imageio (or at least based on)
  56. import os
  57. import time
  58. try:
  59. import PIL
  60. from PIL import Image
  61. pillow = True
  62. try:
  63. from PIL import PILLOW_VERSION # test if user has Pillow or PIL
  64. except ImportError:
  65. pillow = False
  66. from PIL.GifImagePlugin import getheader, getdata
  67. except ImportError:
  68. PIL = None
  69. try:
  70. import numpy as np
  71. except ImportError:
  72. np = None
  73. def get_cKDTree():
  74. try:
  75. from scipy.spatial import cKDTree
  76. except ImportError:
  77. cKDTree = None
  78. return cKDTree
  79. # getheader gives a 87a header and a color palette (two elements in a list)
  80. # getdata()[0] gives the Image Descriptor up to (including) "LZW min code size"
  81. # getdatas()[1:] is the image data itself in chuncks of 256 bytes (well
  82. # technically the first byte says how many bytes follow, after which that
  83. # amount (max 255) follows)
  84. def checkImages(images):
  85. """ checkImages(images)
  86. Check numpy images and correct intensity range etc.
  87. The same for all movie formats.
  88. :param images:
  89. """
  90. # Init results
  91. images2 = []
  92. for im in images:
  93. if PIL and isinstance(im, PIL.Image.Image):
  94. # We assume PIL images are allright
  95. images2.append(im)
  96. elif np and isinstance(im, np.ndarray):
  97. # Check and convert dtype
  98. if im.dtype == np.uint8:
  99. images2.append(im) # Ok
  100. elif im.dtype in [np.float32, np.float64]:
  101. im = im.copy()
  102. im[im < 0] = 0
  103. im[im > 1] = 1
  104. im *= 255
  105. images2.append(im.astype(np.uint8))
  106. else:
  107. im = im.astype(np.uint8)
  108. images2.append(im)
  109. # Check size
  110. if im.ndim == 2:
  111. pass # ok
  112. elif im.ndim == 3:
  113. if im.shape[2] not in [3, 4]:
  114. raise ValueError('This array can not represent an image.')
  115. else:
  116. raise ValueError('This array can not represent an image.')
  117. else:
  118. raise ValueError('Invalid image type: ' + str(type(im)))
  119. # Done
  120. return images2
  121. def intToBin(i):
  122. """Integer to two bytes"""
  123. # devide in two parts (bytes)
  124. i1 = i % 256
  125. i2 = int(i / 256)
  126. # make string (little endian)
  127. return chr(i1) + chr(i2)
  128. class GifWriter:
  129. """Class that contains methods for helping write the animated GIF file.
  130. """
  131. def getheaderAnim(self, im):
  132. """Get animation header. To replace PILs getheader()[0]
  133. :param im:
  134. """
  135. bb = "GIF89a"
  136. bb += intToBin(im.size[0])
  137. bb += intToBin(im.size[1])
  138. bb += "\x87\x00\x00"
  139. return bb
  140. def getImageDescriptor(self, im, xy=None):
  141. """Used for the local color table properties per image.
  142. Otherwise global color table applies to all frames irrespective of
  143. whether additional colors comes in play that require a redefined
  144. palette. Still a maximum of 256 color per frame, obviously.
  145. Written by Ant1 on 2010-08-22
  146. Modified by Alex Robinson in Janurari 2011 to implement subrectangles.
  147. :param im:
  148. :param xy:
  149. """
  150. # Defaule use full image and place at upper left
  151. if xy is None:
  152. xy = (0, 0)
  153. # Image separator,
  154. bb = '\x2C'
  155. # Image position and size
  156. bb += intToBin(xy[0]) # Left position
  157. bb += intToBin(xy[1]) # Top position
  158. bb += intToBin(im.size[0]) # image width
  159. bb += intToBin(im.size[1]) # image height
  160. # packed field: local color table flag1, interlace0, sorted table0,
  161. # reserved00, lct size111=7=2^(7 + 1)=256.
  162. bb += '\x87'
  163. # LZW min size code now comes later, begining of [image data] blocks
  164. return bb
  165. def getAppExt(self, loops=float('inf')):
  166. """Application extention. This part specifies the amount of loops.
  167. If loops is 0 or inf, it goes on infinitely.
  168. :param float loops:
  169. """
  170. if loops == 0 or loops == float('inf'):
  171. loops = 2 ** 16 - 1
  172. #bb = "" # application extension should not be used
  173. # (the extension interprets zero loops
  174. # to mean an infinite number of loops)
  175. # Mmm, does not seem to work
  176. if True:
  177. bb = "\x21\xFF\x0B" # application extension
  178. bb += "NETSCAPE2.0"
  179. bb += "\x03\x01"
  180. bb += intToBin(loops)
  181. bb += '\x00' # end
  182. return bb
  183. def getGraphicsControlExt(self, duration=0.1, dispose=2):
  184. """Graphics Control Extension. A sort of header at the start of
  185. each image. Specifies duration and transparancy.
  186. Dispose
  187. -------
  188. * 0 - No disposal specified.
  189. * 1 - Do not dispose. The graphic is to be left in place.
  190. * 2 - Restore to background color. The area used by the graphic
  191. must be restored to the background color.
  192. * 3 - Restore to previous. The decoder is required to restore the
  193. area overwritten by the graphic with what was there prior to
  194. rendering the graphic.
  195. * 4-7 -To be defined.
  196. :param double duration:
  197. :param dispose:
  198. """
  199. bb = '\x21\xF9\x04'
  200. bb += chr((dispose & 3) << 2) # low bit 1 == transparency,
  201. # 2nd bit 1 == user input , next 3 bits, the low two of which are used,
  202. # are dispose.
  203. bb += intToBin(int(duration * 100)) # in 100th of seconds
  204. bb += '\x00' # no transparant color
  205. bb += '\x00' # end
  206. return bb
  207. def handleSubRectangles(self, images, subRectangles):
  208. """Handle the sub-rectangle stuff. If the rectangles are given by the
  209. user, the values are checked. Otherwise the subrectangles are
  210. calculated automatically.
  211. """
  212. if isinstance(subRectangles, (tuple, list)):
  213. # xy given directly
  214. # Check xy
  215. xy = subRectangles
  216. if xy is None:
  217. xy = (0, 0)
  218. if hasattr(xy, '__len__'):
  219. if len(xy) == len(images):
  220. xy = [xxyy for xxyy in xy]
  221. else:
  222. raise ValueError("len(xy) doesn't match amount of images.")
  223. else:
  224. xy = [xy for im in images]
  225. xy[0] = (0, 0)
  226. else:
  227. # Calculate xy using some basic image processing
  228. # Check Numpy
  229. if np is None:
  230. raise RuntimeError("Need Numpy to use auto-subRectangles.")
  231. # First make numpy arrays if required
  232. for i in range(len(images)):
  233. im = images[i]
  234. if isinstance(im, Image.Image):
  235. tmp = im.convert() # Make without palette
  236. a = np.asarray(tmp)
  237. if len(a.shape) == 0:
  238. raise MemoryError("Too little memory to convert PIL image to array")
  239. images[i] = a
  240. # Determine the sub rectangles
  241. images, xy = self.getSubRectangles(images)
  242. # Done
  243. return images, xy
  244. def getSubRectangles(self, ims):
  245. """ getSubRectangles(ims)
  246. Calculate the minimal rectangles that need updating each frame.
  247. Returns a two-element tuple containing the cropped images and a
  248. list of x-y positions.
  249. Calculating the subrectangles takes extra time, obviously. However,
  250. if the image sizes were reduced, the actual writing of the GIF
  251. goes faster. In some cases applying this method produces a GIF faster.
  252. """
  253. # Check image count
  254. if len(ims) < 2:
  255. return ims, [(0, 0) for i in ims]
  256. # We need numpy
  257. if np is None:
  258. raise RuntimeError("Need Numpy to calculate sub-rectangles. ")
  259. # Prepare
  260. ims2 = [ims[0]]
  261. xy = [(0, 0)]
  262. t0 = time.time()
  263. # Iterate over images
  264. prev = ims[0]
  265. for im in ims[1:]:
  266. # Get difference, sum over colors
  267. diff = np.abs(im-prev)
  268. if diff.ndim == 3:
  269. diff = diff.sum(2)
  270. # Get begin and end for both dimensions
  271. X = np.argwhere(diff.sum(0))
  272. Y = np.argwhere(diff.sum(1))
  273. # Get rect coordinates
  274. if X.size and Y.size:
  275. x0, x1 = X[0], X[-1] + 1
  276. y0, y1 = Y[0], Y[-1] + 1
  277. else: # No change ... make it minimal
  278. x0, x1 = 0, 2
  279. y0, y1 = 0, 2
  280. # Cut out and store
  281. im2 = im[y0:y1, x0:x1]
  282. prev = im
  283. ims2.append(im2)
  284. xy.append((x0, y0))
  285. # Done
  286. # print('%1.2f seconds to determine subrectangles of %i images' %
  287. # (time.time()-t0, len(ims2)))
  288. return ims2, xy
  289. def convertImagesToPIL(self, images, dither, nq=0):
  290. """ convertImagesToPIL(images, nq=0)
  291. Convert images to Paletted PIL images, which can then be
  292. written to a single animaged GIF.
  293. """
  294. # Convert to PIL images
  295. images2 = []
  296. for im in images:
  297. if isinstance(im, Image.Image):
  298. images2.append(im)
  299. elif np and isinstance(im, np.ndarray):
  300. if im.ndim == 3 and im.shape[2] == 3:
  301. im = Image.fromarray(im, 'RGB')
  302. elif im.ndim == 3 and im.shape[2] == 4:
  303. im = Image.fromarray(im[:, :, :3], 'RGB')
  304. elif im.ndim == 2:
  305. im = Image.fromarray(im, 'L')
  306. images2.append(im)
  307. # Convert to paletted PIL images
  308. images, images2 = images2, []
  309. if nq >= 1:
  310. # NeuQuant algorithm
  311. for im in images:
  312. im = im.convert("RGBA") # NQ assumes RGBA
  313. nqInstance = NeuQuant(im, int(nq)) # Learn colors from image
  314. if dither:
  315. im = im.convert("RGB").quantize(palette=nqInstance.paletteImage())
  316. else:
  317. # Use to quantize the image itself
  318. im = nqInstance.quantize(im)
  319. images2.append(im)
  320. else:
  321. # Adaptive PIL algorithm
  322. AD = Image.ADAPTIVE
  323. for im in images:
  324. im = im.convert('P', palette=AD, dither=dither)
  325. images2.append(im)
  326. # Done
  327. return images2
  328. def writeGifToFile(self, fp, images, durations, loops, xys, disposes):
  329. """ writeGifToFile(fp, images, durations, loops, xys, disposes)
  330. Given a set of images writes the bytes to the specified stream.
  331. Requires different handling of palette for PIL and Pillow:
  332. based on https://github.com/rec/echomesh/blob/master/
  333. code/python/external/images2gif.py
  334. """
  335. # Obtain palette for all images and count each occurance
  336. palettes, occur = [], []
  337. for im in images:
  338. if not pillow:
  339. palette = getheader(im)[1]
  340. else:
  341. palette = getheader(im)[0][-1]
  342. if not palette:
  343. palette = im.palette.tobytes()
  344. palettes.append(palette)
  345. for palette in palettes:
  346. occur.append(palettes.count(palette))
  347. # Select most-used palette as the global one (or first in case no max)
  348. globalPalette = palettes[occur.index(max(occur))]
  349. # Init
  350. frames = 0
  351. firstFrame = True
  352. for im, palette in zip(images, palettes):
  353. if firstFrame:
  354. # Write header
  355. # Gather info
  356. header = self.getheaderAnim(im)
  357. appext = self.getAppExt(loops)
  358. # Write
  359. fp.write(header)
  360. fp.write(globalPalette)
  361. fp.write(appext)
  362. # Next frame is not the first
  363. firstFrame = False
  364. if True:
  365. # Write palette and image data
  366. # Gather info
  367. data = getdata(im)
  368. imdes, data = data[0], data[1:]
  369. graphext = self.getGraphicsControlExt(durations[frames],
  370. disposes[frames])
  371. # Make image descriptor suitable for using 256 local color palette
  372. lid = self.getImageDescriptor(im, xys[frames])
  373. # Write local header
  374. if (palette != globalPalette) or (disposes[frames] != 2):
  375. # Use local color palette
  376. fp.write(graphext)
  377. fp.write(lid) # write suitable image descriptor
  378. fp.write(palette) # write local color table
  379. fp.write('\x08') # LZW minimum size code
  380. else:
  381. # Use global color palette
  382. fp.write(graphext)
  383. fp.write(imdes) # write suitable image descriptor
  384. # Write image data
  385. for d in data:
  386. fp.write(d)
  387. # Prepare for next round
  388. frames = frames + 1
  389. fp.write(";") # end gif
  390. return frames
  391. ## Exposed functions
  392. def writeGif(filename, images, duration=0.1, repeat=True, dither=False,
  393. nq=0, subRectangles=True, dispose=None):
  394. """Write an animated gif from the specified images.
  395. Parameters
  396. ----------
  397. filename : string
  398. The name of the file to write the image to.
  399. images : list
  400. Should be a list consisting of PIL images or numpy arrays.
  401. The latter should be between 0 and 255 for integer types, and
  402. between 0 and 1 for float types.
  403. duration : scalar or list of scalars
  404. The duration for all frames, or (if a list) for each frame.
  405. repeat : bool or integer
  406. The amount of loops. If True, loops infinitetely.
  407. dither : bool
  408. Whether to apply dithering
  409. nq : integer
  410. If nonzero, applies the NeuQuant quantization algorithm to create
  411. the color palette. This algorithm is superior, but slower than
  412. the standard PIL algorithm. The value of nq is the quality
  413. parameter. 1 represents the best quality. 10 is in general a
  414. good tradeoff between quality and speed. When using this option,
  415. better results are usually obtained when subRectangles is False.
  416. subRectangles : False, True, or a list of 2-element tuples
  417. Whether to use sub-rectangles. If True, the minimal rectangle that
  418. is required to update each frame is automatically detected. This
  419. can give significant reductions in file size, particularly if only
  420. a part of the image changes. One can also give a list of x-y
  421. coordinates if you want to do the cropping yourself. The default
  422. is True.
  423. dispose : int
  424. How to dispose each frame. 1 means that each frame is to be left
  425. in place. 2 means the background color should be restored after
  426. each frame. 3 means the decoder should restore the previous frame.
  427. If subRectangles==False, the default is 2, otherwise it is 1.
  428. """
  429. # Check PIL
  430. if PIL is None:
  431. raise RuntimeError("Need PIL to write animated gif files.")
  432. # Check images
  433. images = checkImages(images)
  434. # Instantiate writer object
  435. gifWriter = GifWriter()
  436. # Check loops
  437. if repeat is False:
  438. loops = 1
  439. elif repeat is True:
  440. loops = 0 # zero means infinite
  441. else:
  442. loops = int(repeat)
  443. # Check duration
  444. if hasattr(duration, '__len__'):
  445. if len(duration) == len(images):
  446. duration = [d for d in duration]
  447. else:
  448. raise ValueError("len(duration) doesn't match amount of images.")
  449. else:
  450. duration = [duration for im in images]
  451. # Check subrectangles
  452. if subRectangles:
  453. images, xy = gifWriter.handleSubRectangles(images, subRectangles)
  454. defaultDispose = 1 # Leave image in place
  455. else:
  456. # Normal mode
  457. xy = [(0, 0) for im in images]
  458. defaultDispose = 2 # Restore to background color.
  459. # Check dispose
  460. if dispose is None:
  461. dispose = defaultDispose
  462. if hasattr(dispose, '__len__'):
  463. if len(dispose) != len(images):
  464. raise ValueError("len(xy) doesn't match amount of images.")
  465. else:
  466. dispose = [dispose for im in images]
  467. # Make images in a format that we can write easy
  468. images = gifWriter.convertImagesToPIL(images, dither, nq)
  469. # Write
  470. fp = open(filename, 'wb')
  471. try:
  472. gifWriter.writeGifToFile(fp, images, duration, loops, xy, dispose)
  473. finally:
  474. fp.close()
  475. def readGif(filename, asNumpy=True):
  476. """Read images from an animated GIF file. Returns a list of numpy
  477. arrays, or, if asNumpy is false, a list if PIL images.
  478. """
  479. # Check PIL
  480. if PIL is None:
  481. raise RuntimeError("Need PIL to read animated gif files.")
  482. # Check Numpy
  483. if np is None:
  484. raise RuntimeError("Need Numpy to read animated gif files.")
  485. # Check whether it exists
  486. if not os.path.isfile(filename):
  487. raise IOError('File not found: ' + str(filename))
  488. # Load file using PIL
  489. pilIm = PIL.Image.open(filename)
  490. pilIm.seek(0)
  491. # Read all images inside
  492. images = []
  493. try:
  494. while True:
  495. # Get image as numpy array
  496. tmp = pilIm.convert() # Make without palette
  497. a = np.asarray(tmp)
  498. if len(a.shape) == 0:
  499. raise MemoryError("Too little memory to convert PIL image to array")
  500. # Store, and next
  501. images.append(a)
  502. pilIm.seek(pilIm.tell() + 1)
  503. except EOFError:
  504. pass
  505. # Convert to normal PIL images if needed
  506. if not asNumpy:
  507. images2 = images
  508. images = []
  509. for im in images2:
  510. images.append(PIL.Image.fromarray(im))
  511. # Done
  512. return images
  513. class NeuQuant:
  514. """ NeuQuant(image, samplefac=10, colors=256)
  515. samplefac should be an integer number of 1 or higher, 1
  516. being the highest quality, but the slowest performance.
  517. With avalue of 10, one tenth of all pixels are used during
  518. training. This value seems a nice tradeof between speed
  519. and quality.
  520. colors is the amount of colors to reduce the image to. This
  521. should best be a power of two.
  522. See also:
  523. http://members.ozemail.com.au/~dekker/NEUQUANT.HTML
  524. License of the NeuQuant Neural-Net Quantization Algorithm
  525. ---------------------------------------------------------
  526. Copyright (c) 1994 Anthony Dekker
  527. Ported to python by Marius van Voorden in 2010
  528. NEUQUANT Neural-Net quantization algorithm by Anthony Dekker, 1994.
  529. See "Kohonen neural networks for optimal colour quantization"
  530. in "network: Computation in Neural Systems" Vol. 5 (1994) pp 351-367.
  531. for a discussion of the algorithm.
  532. See also http://members.ozemail.com.au/~dekker/NEUQUANT.HTML
  533. Any party obtaining a copy of these files from the author, directly or
  534. indirectly, is granted, free of charge, a full and unrestricted
  535. irrevocable, world-wide, paid up, royalty-free, nonexclusive right and
  536. license to deal in this software and documentation files (the "Software"),
  537. including without limitation the rights to use, copy, modify, merge,
  538. publish, distribute, sublicense, and/or sell copies of the Software, and
  539. to permit persons who receive copies from any such party to do so, with
  540. the only requirement being that this copyright notice remain intact.
  541. """
  542. NCYCLES = None # Number of learning cycles
  543. NETSIZE = None # Number of colours used
  544. SPECIALS = None # Number of reserved colours used
  545. BGCOLOR = None # Reserved background colour
  546. CUTNETSIZE = None
  547. MAXNETPOS = None
  548. INITRAD = None # For 256 colours, radius starts at 32
  549. RADIUSBIASSHIFT = None
  550. RADIUSBIAS = None
  551. INITBIASRADIUS = None
  552. RADIUSDEC = None # Factor of 1/30 each cycle
  553. ALPHABIASSHIFT = None
  554. INITALPHA = None # biased by 10 bits
  555. GAMMA = None
  556. BETA = None
  557. BETAGAMMA = None
  558. network = None # The network itself
  559. colormap = None # The network itself
  560. netindex = None # For network lookup - really 256
  561. bias = None # Bias and freq arrays for learning
  562. freq = None
  563. pimage = None
  564. # Four primes near 500 - assume no image has a length so large
  565. # that it is divisible by all four primes
  566. PRIME1 = 499
  567. PRIME2 = 491
  568. PRIME3 = 487
  569. PRIME4 = 503
  570. MAXPRIME = PRIME4
  571. pixels = None
  572. samplefac = None
  573. a_s = None
  574. def setconstants(self, samplefac, colors):
  575. self.NCYCLES = 100 # Number of learning cycles
  576. self.NETSIZE = colors # Number of colours used
  577. self.SPECIALS = 3 # Number of reserved colours used
  578. self.BGCOLOR = self.SPECIALS-1 # Reserved background colour
  579. self.CUTNETSIZE = self.NETSIZE - self.SPECIALS
  580. self.MAXNETPOS = self.NETSIZE - 1
  581. self.INITRAD = self.NETSIZE/8 # For 256 colours, radius starts at 32
  582. self.RADIUSBIASSHIFT = 6
  583. self.RADIUSBIAS = 1 << self.RADIUSBIASSHIFT
  584. self.INITBIASRADIUS = self.INITRAD * self.RADIUSBIAS
  585. self.RADIUSDEC = 30 # Factor of 1/30 each cycle
  586. self.ALPHABIASSHIFT = 10 # Alpha starts at 1
  587. self.INITALPHA = 1 << self.ALPHABIASSHIFT # biased by 10 bits
  588. self.GAMMA = 1024.0
  589. self.BETA = 1.0/1024.0
  590. self.BETAGAMMA = self.BETA * self.GAMMA
  591. self.network = np.empty((self.NETSIZE, 3), dtype='float64') # The network itself
  592. self.colormap = np.empty((self.NETSIZE, 4), dtype='int32') # The network itself
  593. self.netindex = np.empty(256, dtype='int32') # For network lookup - really 256
  594. self.bias = np.empty(self.NETSIZE, dtype='float64') # Bias and freq arrays for learning
  595. self.freq = np.empty(self.NETSIZE, dtype='float64')
  596. self.pixels = None
  597. self.samplefac = samplefac
  598. self.a_s = {}
  599. def __init__(self, image, samplefac=10, colors=256):
  600. # Check Numpy
  601. if np is None:
  602. raise RuntimeError("Need Numpy for the NeuQuant algorithm.")
  603. # Check image
  604. if image.size[0] * image.size[1] < NeuQuant.MAXPRIME:
  605. raise IOError("Image is too small")
  606. if image.mode != "RGBA":
  607. raise IOError("Image mode should be RGBA.")
  608. # Initialize
  609. self.setconstants(samplefac, colors)
  610. self.pixels = np.fromstring(image.tostring(), np.uint32)
  611. self.setUpArrays()
  612. self.learn()
  613. self.fix()
  614. self.inxbuild()
  615. def writeColourMap(self, rgb, outstream):
  616. for i in range(self.NETSIZE):
  617. bb = self.colormap[i, 0]
  618. gg = self.colormap[i, 1]
  619. rr = self.colormap[i, 2]
  620. outstream.write(rr if rgb else bb)
  621. outstream.write(gg)
  622. outstream.write(bb if rgb else rr)
  623. return self.NETSIZE
  624. def setUpArrays(self):
  625. self.network[0, 0] = 0.0 # Black
  626. self.network[0, 1] = 0.0
  627. self.network[0, 2] = 0.0
  628. self.network[1, 0] = 255.0 # White
  629. self.network[1, 1] = 255.0
  630. self.network[1, 2] = 255.0
  631. # RESERVED self.BGCOLOR # Background
  632. for i in range(self.SPECIALS):
  633. self.freq[i] = 1.0 / self.NETSIZE
  634. self.bias[i] = 0.0
  635. for i in range(self.SPECIALS, self.NETSIZE):
  636. p = self.network[i]
  637. p[:] = (255.0 * (i-self.SPECIALS)) / self.CUTNETSIZE
  638. self.freq[i] = 1.0 / self.NETSIZE
  639. self.bias[i] = 0.0
  640. # Omitted: setPixels
  641. def altersingle(self, alpha, i, b, g, r):
  642. """Move neuron i towards biased (b, g, r) by factor alpha"""
  643. n = self.network[i] # Alter hit neuron
  644. n[0] -= (alpha * (n[0] - b))
  645. n[1] -= (alpha * (n[1] - g))
  646. n[2] -= (alpha * (n[2] - r))
  647. def geta(self, alpha, rad):
  648. try:
  649. return self.a_s[(alpha, rad)]
  650. except KeyError:
  651. length = rad * 2-1
  652. mid = length/2
  653. q = np.array(list(range(mid-1, -1, -1)) + list(range(-1, mid)))
  654. a = alpha * (rad * rad - q * q)/(rad * rad)
  655. a[mid] = 0
  656. self.a_s[(alpha, rad)] = a
  657. return a
  658. def alterneigh(self, alpha, rad, i, b, g, r):
  659. if i-rad >= self.SPECIALS-1:
  660. lo = i-rad
  661. start = 0
  662. else:
  663. lo = self.SPECIALS-1
  664. start = (self.SPECIALS-1 - (i-rad))
  665. if i + rad <= self.NETSIZE:
  666. hi = i + rad
  667. end = rad * 2-1
  668. else:
  669. hi = self.NETSIZE
  670. end = (self.NETSIZE - (i + rad))
  671. a = self.geta(alpha, rad)[start:end]
  672. p = self.network[lo + 1:hi]
  673. p -= np.transpose(np.transpose(p - np.array([b, g, r])) * a)
  674. #def contest(self, b, g, r):
  675. # """ Search for biased BGR values
  676. # Finds closest neuron (min dist) and updates self.freq
  677. # finds best neuron (min dist-self.bias) and returns position
  678. # for frequently chosen neurons, self.freq[i] is high and self.bias[i] is negative
  679. # self.bias[i] = self.GAMMA * ((1/self.NETSIZE)-self.freq[i])"""
  680. #
  681. # i, j = self.SPECIALS, self.NETSIZE
  682. # dists = abs(self.network[i:j] - np.array([b, g, r])).sum(1)
  683. # bestpos = i + np.argmin(dists)
  684. # biasdists = dists - self.bias[i:j]
  685. # bestbiaspos = i + np.argmin(biasdists)
  686. # self.freq[i:j] -= self.BETA * self.freq[i:j]
  687. # self.bias[i:j] += self.BETAGAMMA * self.freq[i:j]
  688. # self.freq[bestpos] += self.BETA
  689. # self.bias[bestpos] -= self.BETAGAMMA
  690. # return bestbiaspos
  691. def contest(self, b, g, r):
  692. """Search for biased BGR values
  693. Finds closest neuron (min dist) and updates self.freq
  694. finds best neuron (min dist-self.bias) and returns position
  695. for frequently chosen neurons, self.freq[i] is high and self.bias[i]
  696. is negative self.bias[i] = self.GAMMA * ((1/self.NETSIZE)-self.freq[i])
  697. """
  698. i, j = self.SPECIALS, self.NETSIZE
  699. dists = abs(self.network[i:j] - np.array([b, g, r])).sum(1)
  700. bestpos = i + np.argmin(dists)
  701. biasdists = dists - self.bias[i:j]
  702. bestbiaspos = i + np.argmin(biasdists)
  703. self.freq[i:j] *= (1-self.BETA)
  704. self.bias[i:j] += self.BETAGAMMA * self.freq[i:j]
  705. self.freq[bestpos] += self.BETA
  706. self.bias[bestpos] -= self.BETAGAMMA
  707. return bestbiaspos
  708. def specialFind(self, b, g, r):
  709. for i in range(self.SPECIALS):
  710. n = self.network[i]
  711. if n[0] == b and n[1] == g and n[2] == r:
  712. return i
  713. return -1
  714. def learn(self):
  715. biasRadius = self.INITBIASRADIUS
  716. alphadec = 30 + ((self.samplefac-1)/3)
  717. lengthcount = self.pixels.size
  718. samplepixels = lengthcount / self.samplefac
  719. delta = samplepixels / self.NCYCLES
  720. alpha = self.INITALPHA
  721. i = 0
  722. rad = biasRadius >> self.RADIUSBIASSHIFT
  723. if rad <= 1:
  724. rad = 0
  725. print("Beginning 1D learning: samplepixels = %1.2f rad = %i" %
  726. (samplepixels, rad))
  727. step = 0
  728. pos = 0
  729. if lengthcount % NeuQuant.PRIME1 != 0:
  730. step = NeuQuant.PRIME1
  731. elif lengthcount % NeuQuant.PRIME2 != 0:
  732. step = NeuQuant.PRIME2
  733. elif lengthcount % NeuQuant.PRIME3 != 0:
  734. step = NeuQuant.PRIME3
  735. else:
  736. step = NeuQuant.PRIME4
  737. i = 0
  738. printed_string = ''
  739. while i < samplepixels:
  740. if i % 100 == 99:
  741. tmp = '\b' * len(printed_string)
  742. printed_string = str((i + 1) * 100/samplepixels) + "%\n"
  743. print(tmp + printed_string)
  744. p = self.pixels[pos]
  745. r = (p >> 16) & 0xff
  746. g = (p >> 8) & 0xff
  747. b = (p) & 0xff
  748. if i == 0: # Remember background colour
  749. self.network[self.BGCOLOR] = [b, g, r]
  750. j = self.specialFind(b, g, r)
  751. if j < 0:
  752. j = self.contest(b, g, r)
  753. if j >= self.SPECIALS: # Don't learn for specials
  754. a = (1.0 * alpha) / self.INITALPHA
  755. self.altersingle(a, j, b, g, r)
  756. if rad > 0:
  757. self.alterneigh(a, rad, j, b, g, r)
  758. pos = (pos + step) % lengthcount
  759. i += 1
  760. if i % delta == 0:
  761. alpha -= alpha / alphadec
  762. biasRadius -= biasRadius / self.RADIUSDEC
  763. rad = biasRadius >> self.RADIUSBIASSHIFT
  764. if rad <= 1:
  765. rad = 0
  766. finalAlpha = (1.0 * alpha)/self.INITALPHA
  767. print("Finished 1D learning: final alpha = %1.2f!" % finalAlpha)
  768. def fix(self):
  769. for i in range(self.NETSIZE):
  770. for j in range(3):
  771. x = int(0.5 + self.network[i, j])
  772. x = max(0, x)
  773. x = min(255, x)
  774. self.colormap[i, j] = x
  775. self.colormap[i, 3] = i
  776. def inxbuild(self):
  777. previouscol = 0
  778. startpos = 0
  779. for i in range(self.NETSIZE):
  780. p = self.colormap[i]
  781. q = None
  782. smallpos = i
  783. smallval = p[1] # Index on g
  784. # Find smallest in i..self.NETSIZE-1
  785. for j in range(i + 1, self.NETSIZE):
  786. q = self.colormap[j]
  787. if q[1] < smallval: # Index on g
  788. smallpos = j
  789. smallval = q[1] # Index on g
  790. q = self.colormap[smallpos]
  791. # Swap p (i) and q (smallpos) entries
  792. if i != smallpos:
  793. p[:], q[:] = q, p.copy()
  794. # smallval entry is now in position i
  795. if smallval != previouscol:
  796. self.netindex[previouscol] = (startpos + i) >> 1
  797. for j in range(previouscol + 1, smallval):
  798. self.netindex[j] = i
  799. previouscol = smallval
  800. startpos = i
  801. self.netindex[previouscol] = (startpos + self.MAXNETPOS) >> 1
  802. for j in range(previouscol + 1, 256): # Really 256
  803. self.netindex[j] = self.MAXNETPOS
  804. def paletteImage(self):
  805. """PIL weird interface for making a paletted image: create an image
  806. which already has the palette, and use that in Image.quantize. This
  807. function returns this palette image."""
  808. if self.pimage is None:
  809. palette = []
  810. for i in range(self.NETSIZE):
  811. palette.extend(self.colormap[i][:3])
  812. palette.extend([0] * (256-self.NETSIZE) * 3)
  813. # a palette image to use for quant
  814. self.pimage = Image.new("P", (1, 1), 0)
  815. self.pimage.putpalette(palette)
  816. return self.pimage
  817. def quantize(self, image):
  818. """Use a kdtree to quickly find the closest palette colors for the
  819. pixels
  820. :param image:
  821. """
  822. if get_cKDTree():
  823. return self.quantize_with_scipy(image)
  824. else:
  825. print('Scipy not available, falling back to slower version.')
  826. return self.quantize_without_scipy(image)
  827. def quantize_with_scipy(self, image):
  828. w, h = image.size
  829. px = np.asarray(image).copy()
  830. px2 = px[:, :, :3].reshape((w * h, 3))
  831. cKDTree = get_cKDTree()
  832. kdtree = cKDTree(self.colormap[:, :3], leafsize=10)
  833. result = kdtree.query(px2)
  834. colorindex = result[1]
  835. print("Distance: %1.2f" % (result[0].sum()/(w * h)))
  836. px2[:] = self.colormap[colorindex, :3]
  837. return Image.fromarray(px).convert("RGB").quantize(palette=self.paletteImage())
  838. def quantize_without_scipy(self, image):
  839. """" This function can be used if no scipy is availabe.
  840. It's 7 times slower though.
  841. :param image:
  842. """
  843. w, h = image.size
  844. px = np.asarray(image).copy()
  845. memo = {}
  846. for j in range(w):
  847. for i in range(h):
  848. key = (px[i, j, 0], px[i, j, 1], px[i, j, 2])
  849. try:
  850. val = memo[key]
  851. except KeyError:
  852. val = self.convert(*key)
  853. memo[key] = val
  854. px[i, j, 0], px[i, j, 1], px[i, j, 2] = val
  855. return Image.fromarray(px).convert("RGB").quantize(palette=self.paletteImage())
  856. def convert(self, *color):
  857. i = self.inxsearch(*color)
  858. return self.colormap[i, :3]
  859. def inxsearch(self, r, g, b):
  860. """Search for BGR values 0..255 and return colour index"""
  861. dists = (self.colormap[:, :3] - np.array([r, g, b]))
  862. a = np.argmin((dists * dists).sum(1))
  863. return a
  864. if __name__ == '__main__':
  865. im = np.zeros((200, 200), dtype=np.uint8)
  866. im[10: 30, :] = 100
  867. im[:, 80: 120] = 255
  868. im[-50: -40, :] = 50
  869. images = [im * 1.0, im * 0.8, im * 0.6, im * 0.4, im * 0]
  870. writeGif('lala3.gif', images, duration=0.5, dither=0)