1  """
 2  ImageHash module tests
 3  """
 4  
 5  import unittest
 6  
 7  from PIL import Image
 8  
 9  from txtai.pipeline import ImageHash
10  
11  # pylint: disable=C0411
12  from utils import Utils
13  
14  
15  class TestImageHash(unittest.TestCase):
16      """
17      ImageHash tests.
18      """
19  
20      @classmethod
21      def setUpClass(cls):
22          """
23          Caches an image to hash
24          """
25  
26          cls.image = Image.open(Utils.PATH + "/books.jpg")
27  
28      def testArray(self):
29          """
30          Test numpy return type
31          """
32  
33          ihash = ImageHash(strings=False)
34          self.assertEqual(ihash(self.image).shape, (64,))
35  
36      def testAverage(self):
37          """
38          Test average hash
39          """
40  
41          ihash = ImageHash("average")
42          self.assertIn(ihash(self.image), ["0859dd04bfbfbf00", "0859dd04ffbfbf00"])
43  
44      def testColor(self):
45          """
46          Test color hash
47          """
48  
49          ihash = ImageHash("color")
50          self.assertIn(ihash(self.image), ["1ffffe02000e000c0e0000070000", "1ff8fe03000e00070e0000070000"])
51  
52      def testDifference(self):
53          """
54          Test difference hash
55          """
56  
57          ihash = ImageHash("difference")
58          self.assertEqual(ihash(self.image), "d291996d6969686a")
59  
60      def testPerceptual(self):
61          """
62          Test perceptual hash
63          """
64  
65          ihash = ImageHash("perceptual")
66          self.assertEqual(ihash(self.image), "8be8418577b331b9")
67  
68      def testWavelet(self):
69          """
70          Test wavelet hash
71          """
72  
73          ihash = ImageHash("wavelet")
74          self.assertEqual(ihash(Utils.PATH + "/books.jpg"), "68015d85bfbf3f00")