add support for mexican hat wavelet in pHash

westonplatter · Mar 1, 2014 · fa0ed12 · fa0ed12
1 parent 2ace3be
commit fa0ed12
Show file tree

Hide file tree

Showing 3 changed files with 62 additions and 0 deletions.
diff --git a/ext/phashion_ext/phashion_ext.c b/ext/phashion_ext/phashion_ext.c
@@ -54,6 +54,50 @@ static VALUE hamming_distance(VALUE self, VALUE a, VALUE b) {
     return INT2NUM(result);
 }
 
+static VALUE hamming_distance2(VALUE self, VALUE a, VALUE b) {
+	double distance;
+	uint8_t* hashA;
+	uint8_t* hashB;
+	int lenA = RARRAY_LEN(a);
+	int lenB = RARRAY_LEN(b);
+
+	hashA = (uint8_t*)xcalloc(lenA, sizeof(uint8_t));
+	hashB = (uint8_t*)xcalloc(lenB, sizeof(uint8_t));
+
+	for(int i = 0; i < lenA; i++) {
+		hashA[i] = NUM2INT(rb_ary_entry(a, i));
+	}
+	for(int i = 0; i < lenB; i++) {
+		hashB[i] = NUM2INT(rb_ary_entry(b, i));
+	}
+
+	distance = ph_hammingdistance2(hashA, lenA, hashB, lenB);
+
+	xfree(hashA);
+	xfree(hashB);
+
+	return DBL2NUM(distance);
+}
+
+static VALUE mh_hash_for(VALUE self, VALUE filename, VALUE alpha, VALUE lvl) {
+	uint8_t* result;
+	int n;
+	VALUE array;
+	result = ph_mh_imagehash(StringValuePtr(filename),
+			n,
+			NUM2DBL(alpha),
+			NUM2DBL(lvl));
+	array = rb_ary_new2(n);
+
+	for(int i = 0; i < n; i++) {
+		rb_ary_push(array, INT2FIX(result[i]));
+	}
+
+	xfree(result);
+
+	return array;
+}
+
 #ifdef __cplusplus
 extern "C" {
 #endif
@@ -63,6 +107,8 @@ extern "C" {
 
     rb_define_singleton_method(c, "hamming_distance", (VALUE(*)(ANYARGS))hamming_distance, 2);
     rb_define_singleton_method(c, "image_hash_for", (VALUE(*)(ANYARGS))image_hash_for, 1);
+    rb_define_singleton_method(c, "_mh_hash_for", (VALUE(*)(ANYARGS))mh_hash_for, 3);
+    rb_define_singleton_method(c, "hamming_distance2", (VALUE(*)(ANYARGS))hamming_distance2, 2);
   }
 
 #ifdef HAVE_SQLITE3EXT_H

diff --git a/lib/phashion.rb b/lib/phashion.rb
@@ -35,6 +35,10 @@ def fingerprint
     end
   end
 
+  def self.mh_hash_for(filename, alpha = 2.0, lvl = 1.0)
+    _mh_hash_for(filename, alpha, lvl)
+  end
+
   def self.so_file
     extname = RbConfig::CONFIG['DLEXT']
     File.join File.dirname(__FILE__), "phashion_ext.#{extname}"

diff --git a/test/test_phashion.rb b/test/test_phashion.rb
@@ -50,6 +50,18 @@ def test_db_extension
     assert_equal expected, rows.first.first
   end
 
+  def test_mh_hash_for
+    jpg = File.dirname(__FILE__) + '/jpg/Broccoli_Super_Food.jpg'
+    png = File.dirname(__FILE__) + '/png/Broccoli_Super_Food.png'
+
+    hash1 = Phashion.mh_hash_for jpg
+    hash2 = Phashion.mh_hash_for png
+
+    assert_kind_of Array, hash1
+    assert_kind_of Array, hash2
+    assert_in_delta 0.100, Phashion.hamming_distance2(hash1, hash2), 0.0007
+  end
+
   def test_duplicate_detection
     files = %w(86x86-0a1e.jpeg 86x86-83d6.jpeg 86x86-a855.jpeg)
     images = files.map {|f| Phashion::Image.new("#{File.dirname(__FILE__) + '/../test/jpg/'}#{f}")}