You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

248 lines
7.3 KiB

9 years ago
9 years ago
9 years ago
9 years ago
9 years ago
  1. // Copyright 2016 Drew Short <drew@sothr.com>.
  2. //
  3. // Licensed under the MIT license<LICENSE-MIT or http://opensource.org/licenses/MIT>.
  4. // This file may not be copied, modified, or distributed except according to those terms.
  5. extern crate dft;
  6. extern crate image;
  7. use cache::Cache;
  8. use self::image::FilterType;
  9. use std::f64;
  10. use std::path::Path;
  11. use std::fmt;
  12. mod ahash;
  13. mod dhash;
  14. mod phash;
  15. // Constants //
  16. // Used to get ranges for the precision of rounding floats
  17. // Can round to 1 significant factor of precision
  18. const FLOAT_PRECISION_MAX_1: f64 = f64::MAX / 10_f64;
  19. const FLOAT_PRECISION_MIN_1: f64 = f64::MIN / 10_f64;
  20. // Can round to 2 significant factors of precision
  21. const FLOAT_PRECISION_MAX_2: f64 = f64::MAX / 100_f64;
  22. const FLOAT_PRECISION_MIN_2: f64 = f64::MIN / 100_f64;
  23. // Can round to 3 significant factors of precision
  24. const FLOAT_PRECISION_MAX_3: f64 = f64::MAX / 1000_f64;
  25. const FLOAT_PRECISION_MIN_3: f64 = f64::MIN / 1000_f64;
  26. // Can round to 4 significant factors of precision
  27. const FLOAT_PRECISION_MAX_4: f64 = f64::MAX / 10000_f64;
  28. const FLOAT_PRECISION_MIN_4: f64 = f64::MIN / 10000_f64;
  29. // Can round to 5 significant factors of precision
  30. const FLOAT_PRECISION_MAX_5: f64 = f64::MAX / 100000_f64;
  31. const FLOAT_PRECISION_MIN_5: f64 = f64::MIN / 100000_f64;
  32. // Hamming Distance Similarity Limit //
  33. const HAMMING_DISTANCE_SIMILARITY_LIMIT: u64 = 5u64;
  34. // Structs/Enums //
  35. /**
  36. * Prepared image that can be used to generate hashes
  37. */
  38. pub struct PreparedImage<'a> {
  39. orig_path: &'a str,
  40. image: Option<image::DynamicImage>,
  41. }
  42. /**
  43. * Wraps the various perceptual hashes
  44. */
  45. pub struct PerceptualHashes<'a> {
  46. pub orig_path: &'a str,
  47. pub ahash: u64,
  48. pub dhash: u64,
  49. pub phash: u64,
  50. }
  51. impl<'a> PerceptualHashes<'a> {
  52. pub fn similar(&self, other: &'a PerceptualHashes<'a>) -> bool {
  53. if self.orig_path != other.orig_path &&
  54. calculate_hamming_distance(self.ahash, other.ahash) <=
  55. HAMMING_DISTANCE_SIMILARITY_LIMIT &&
  56. calculate_hamming_distance(self.dhash, other.dhash) <=
  57. HAMMING_DISTANCE_SIMILARITY_LIMIT &&
  58. calculate_hamming_distance(self.phash, other.phash) <=
  59. HAMMING_DISTANCE_SIMILARITY_LIMIT {
  60. true
  61. } else {
  62. false
  63. }
  64. }
  65. }
  66. /**
  67. * All the supported precision types
  68. *
  69. * Low aims for 32 bit precision
  70. * Medium aims for 64 bit precision
  71. * High aims for 128 bit precision
  72. */
  73. #[allow(dead_code)]
  74. pub enum Precision {
  75. Low,
  76. Medium,
  77. High,
  78. }
  79. // Get the size of the required image
  80. //
  81. impl Precision {
  82. fn get_size(&self) -> u32 {
  83. match *self {
  84. Precision::Low => 4,
  85. Precision::Medium => 8,
  86. Precision::High => 16,
  87. }
  88. }
  89. }
  90. /**
  91. * Types of hashes supported
  92. */
  93. pub enum HashType {
  94. AHash,
  95. DHash,
  96. PHash,
  97. }
  98. impl fmt::Display for HashType {
  99. fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
  100. match *self {
  101. HashType::AHash => write!(f, "AHash"),
  102. HashType::DHash => write!(f, "DHash"),
  103. HashType::PHash => write!(f, "PHash")
  104. }
  105. }
  106. }
  107. // Traits //
  108. pub trait PerceptualHash {
  109. fn get_hash(&self, cache: &Option<Cache>) -> u64;
  110. }
  111. // Functions //
  112. /**
  113. * Resonsible for parsing a path, converting an image and package it to be
  114. * hashed.
  115. *
  116. * # Arguments
  117. *
  118. * * 'path' - The path to the image requested to be hashed
  119. * * 'size' - The size that the image should be resize to, in the form of size x size
  120. *
  121. * # Returns
  122. *
  123. * A PreparedImage struct with the required information for performing hashing
  124. *
  125. */
  126. pub fn prepare_image<'a>(path: &'a Path,
  127. hash_type: &HashType,
  128. precision: &Precision,
  129. cache: &Option<Cache>)
  130. -> PreparedImage<'a> {
  131. let image_path = path.to_str().unwrap();
  132. let size: u32 = match *hash_type {
  133. HashType::PHash => precision.get_size() * 4,
  134. _ => precision.get_size(),
  135. };
  136. // Check if we have the already converted image in a cache and use that if possible.
  137. match *cache {
  138. Some(ref c) => {
  139. match c.get_image_from_cache(&path, size) {
  140. Some(image) => {
  141. PreparedImage {
  142. orig_path: &*image_path,
  143. image: Some(image),
  144. }
  145. }
  146. None => {
  147. let processed_image = process_image(&image_path, size);
  148. // Oh, and save it in a cache
  149. match processed_image.image {
  150. Some(ref image) => {
  151. match c.put_image_in_cache(&path, size, &image) {
  152. Ok(_) => {}
  153. Err(e) => println!("Unable to store image in cache. {}", e),
  154. };
  155. }
  156. None => {}
  157. };
  158. processed_image
  159. }
  160. }
  161. }
  162. None => process_image(&image_path, size),
  163. }
  164. }
  165. /**
  166. * Turn the image into something we can work with
  167. */
  168. fn process_image<'a>(image_path: &'a str, size: u32) -> PreparedImage<'a> {
  169. // Otherwise let's do that work now and store it.
  170. // println!("Path: {}", image_path);
  171. let image = match image::open(Path::new(image_path)) {
  172. Ok(image) => {
  173. let small_image = image.resize_exact(size, size, FilterType::Lanczos3);
  174. Some(small_image.grayscale())
  175. }
  176. Err(e) => {
  177. println!("Error Processing Image [{}]: {} ", image_path, e);
  178. None
  179. }
  180. };
  181. PreparedImage {
  182. orig_path: &*image_path,
  183. image,
  184. }
  185. }
  186. /**
  187. * Get a specific HashType hash
  188. */
  189. pub fn get_perceptual_hash<'a>(path: &'a Path,
  190. precision: &Precision,
  191. hash_type: &HashType,
  192. cache: &Option<Cache>)
  193. -> u64 {
  194. match *hash_type {
  195. HashType::AHash => ahash::AHash::new(&path, &precision, &cache).get_hash(&cache),
  196. HashType::DHash => dhash::DHash::new(&path, &precision, &cache).get_hash(&cache),
  197. HashType::PHash => phash::PHash::new(&path, &precision, &cache).get_hash(&cache),
  198. }
  199. }
  200. /**
  201. * Get all perceptual hashes for an image
  202. */
  203. pub fn get_perceptual_hashes<'a>(path: &'a Path,
  204. precision: &Precision,
  205. cache: &Option<Cache>)
  206. -> PerceptualHashes<'a> {
  207. let image_path = path.to_str().unwrap();
  208. let ahash = ahash::AHash::new(&path, &precision, &cache).get_hash(&cache);
  209. let dhash = dhash::DHash::new(&path, &precision, &cache).get_hash(&cache);
  210. let phash = phash::PHash::new(&path, &precision, &cache).get_hash(&cache);
  211. PerceptualHashes {
  212. orig_path: &*image_path,
  213. ahash: ahash,
  214. dhash: dhash,
  215. phash: phash,
  216. }
  217. }
  218. /**
  219. * Calculate the number of bits different between two hashes
  220. * Add to the PerceptualHashTrait
  221. */
  222. pub fn calculate_hamming_distance(hash1: u64, hash2: u64) -> u64 {
  223. // The binary xor of the two hashes should give us a number representing
  224. // the differences between the two hashes. All that's left is to count
  225. // the number of 1's in the difference to determine the hamming distance
  226. (hash1 ^ hash2).count_ones() as u64
  227. }