support hashing gif images
[~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 datat type supported by PIL. This
8     disregards mode, depth and meta information. Note that due to limitations
9     in PIL and the image format (interlacing) the full contents are stored and
10     decoded in hexdigest."""
11     maxsize = 1024 * 1024 * 32
12     # max memory usage is about 5 * maxpixels in bytes
13     maxpixels = 1024 * 1024 * 32
14
15     def __init__(self, hashobj):
16         """
17         @param hashobj: a hashlib-like object
18         """
19         self.hashobj = hashobj
20         self.imagedetected = False
21         self.content = io.BytesIO()
22
23     def detect(self):
24         raise NotImplementedError
25
26     def update(self, data):
27         self.content.write(data)
28         if self.content.tell() > self.maxsize:
29             raise ValueError("maximum image size exceeded")
30         if not self.imagedetected:
31             self.imagedetected = self.detect()
32
33     def copy(self):
34         new = self.__class__(self.hashobj.copy())
35         new.imagedetected = self.imagedetected
36         new.content = io.BytesIO(self.content.getvalue())
37         return new
38
39     def hexdigest(self):
40         if not self.imagedetected:
41             raise ValueError("not a image")
42         hashobj = self.hashobj.copy()
43         pos = self.content.tell()
44         try:
45             self.content.seek(0)
46             try:
47                 img = PIL.Image.open(self.content)
48             except IOError:
49                 raise ValueError("broken header")
50             width, height = img.size
51             pack = lambda elem: struct.pack("BBBB", *elem)
52             # special casing easy modes reduces memory usage
53             if img.mode == "L":
54                 pack = lambda elem: struct.pack("BBBB", elem, elem, elem, 255)
55             elif img.mode == "RGB":
56                 pack = lambda elem: struct.pack("BBBB", *(elem + (255,)))
57             elif img.mode != "RGBA":
58                 try:
59                     img = img.convert("RGBA")
60                 except (SyntaxError, IndexError, IOError):
61                     # crazy stuff from PIL
62                     raise ValueError("error reading image")
63             try:
64                 for elem in img.getdata():
65                     hashobj.update(pack(elem))
66             except (SyntaxError, IndexError, IOError): # crazy stuff from PIL
67                 raise ValueError("error reading image")
68         finally:
69             self.content.seek(pos)
70         return "%s%8.8x%8.8x" % (hashobj.hexdigest(), width, height)
71
72
73 class PNGHash(ImageHash):
74     """A hash on the contents of a PNG image."""
75
76     def detect(self):
77         if self.content.tell() < 33: # header + IHDR
78             return False
79         curvalue = self.content.getvalue()
80         if curvalue.startswith(b"\x89PNG\r\n\x1a\n\0\0\0\x0dIHDR"):
81             width, height = struct.unpack(">II", curvalue[16:24])
82             if width * height > self.maxpixels:
83                 raise ValueError("maximum image pixels exceeded")
84             return True
85         raise ValueError("not a png image")
86
87 class GIFHash(ImageHash):
88     """A hash on the contents of the first frame of a GIF image."""
89
90     def detect(self):
91         if self.content.tell() < 10: # magic + logical dimension
92             return False
93         curvalue = self.content.getvalue()
94         if curvalue.startswith((b"GIF87a", "GIF89a")):
95             width, height = struct.unpack("<HH", curvalue[6:10])
96             if width * height > self.maxpixels:
97                 raise ValueError("maximum image pixels exceeded")
98             return True
99         raise ValueError("not a png image")