1148890c9ea167e2037947bf6811655992949ca6
[~helmut/debian-dedup.git] / dedup / image.py
1 import io
2 import struct
3
4 import PIL.Image
5
6 class ImageHash(object):
7     """A hash on the contents of an image. This disregards mode, depth and meta
8     information. Note that due to limitations in PIL and the image format
9     (interlacing) the full contents are stored and decoded in hexdigest."""
10     maxsize = 1024 * 1024 * 32
11     # max memory usage is about 5 * maxpixels in bytes
12     maxpixels = 1024 * 1024 * 32
13
14     def __init__(self, hashobj):
15         """
16         @param hashobj: a hashlib-like object
17         """
18         self.hashobj = hashobj
19         self.imagedetected = False
20         self.content = io.BytesIO()
21
22     def update(self, data):
23         self.content.write(data)
24         if self.content.tell() > self.maxsize:
25             raise ValueError("maximum image size exceeded")
26         if self.imagedetected:
27             return
28         if self.content.tell() < 33: # header + IHDR
29             return
30         curvalue = self.content.getvalue()
31         if curvalue.startswith(b"\x89PNG\r\n\x1a\n\0\0\0\x0dIHDR"):
32             width, height = struct.unpack(">II", curvalue[16:24])
33             if width * height > self.maxpixels:
34                 raise ValueError("maximum image pixels exceeded")
35             self.imagedetected = True
36             return
37         raise ValueError("not a png image")
38
39     def copy(self):
40         new = ImageHash()
41         new.hashobj = self.hashobj.copy()
42         new.imagedetected = self.imagedetected
43         new.content = io.BytesIO(self.content.getvalue())
44         return new
45
46     def hexdigest(self):
47         if not self.imagedetected:
48             raise ValueError("not a png image")
49         hashobj = self.hashobj.copy()
50         pos = self.content.tell()
51         try:
52             self.content.seek(0)
53             try:
54                 img = PIL.Image.open(self.content)
55             except IOError:
56                 raise ValueError("broken png header")
57             width, height = img.size
58             pack = lambda elem: struct.pack("BBBB", *elem)
59             # special casing easy modes reduces memory usage
60             if img.mode == "L":
61                 pack = lambda elem: struct.pack("BBBB", elem, elem, elem, 255)
62             elif img.mode == "RGB":
63                 pack = lambda elem: struct.pack("BBBB", *(elem + (255,)))
64             elif img.mode != "RGBA":
65                 try:
66                     img = img.convert("RGBA")
67                 except (SyntaxError, IndexError, IOError): # crazy stuff from PIL
68                     raise ValueError("error reading png image")
69             try:
70                 for elem in img.getdata():
71                     hashobj.update(pack(elem))
72             except (SyntaxError, IndexError, IOError): # crazy stuff from PIL
73                 raise ValueError("error reading png image")
74         finally:
75             self.content.seek(pos)
76         return "%s%8.8x%8.8x" % (hashobj.hexdigest(), width, height)