source: grass/trunk/lib/python/imaging/images2gif.py

Last change on this file was 71485, checked in by annakrat, 7 years ago

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