remove cairotextoverlay
[numm.git] / numm / image.py
1
2 """
3 Conversion between image files and numpy arrays.
4
5 Images are represented as arrays with shape (height, width, 3).
6 """
7
8 import numpy
9 import Image
10
11 def _get_dimensions(im, width, height):
12     "preserve aspect if width or height is None"
13     if width and height:
14         return (width, height)
15     aspect = im.size[0]/float(im.size[1])
16     if width:
17         return (width, int(width/aspect))
18     else:
19         return (int(height*aspect), height)
20
21 def image2np(path, width=None, height=None):
22     "Load an image file into an array."
23
24     im = Image.open(path)
25
26     if width or height:
27         im = im.resize(_get_dimensions(im, width, height), 1)
28
29     im = im.convert('RGB')
30     np = numpy.asarray(im, dtype=numpy.uint8)
31     return np
32
33 def np2image(np, path):
34     "Save an image array to a file."
35
36     assert np.dtype == numpy.uint8, "nparr must be uint8"
37     if len(np.shape) > 2 and np.shape[2] == 3:
38         mode = 'RGB'
39     else:
40         mode = 'L'
41     im = Image.fromstring(
42         mode, (np.shape[1], np.shape[0]), np.tostring())
43     im.save(path)
44
45 if __name__ == '__main__':
46     # XXX: Replace this with an automated test.
47     import sys
48     np1 = image2np(sys.argv[1])
49     np2image(np1, sys.argv[2])
50     np2 = image2np(sys.argv[2])
51     print 'error: ', numpy.sum(abs(np1-np2))