images2gif.py 36 KB

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