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.

233 lines
8.2 KiB

8 years ago
8 years ago
8 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 image;
  6. use cache::Cache;
  7. use self::image::{GenericImage, DynamicImage};
  8. use std::path::Path;
  9. use super::{HashType, PerceptualHash, Precision, PreparedImage};
  10. use super::dft;
  11. use super::dft::Transform;
  12. use super::image::Pixel;
  13. use super::prepare_image;
  14. pub struct PHash<'a> {
  15. prepared_image: Box<PreparedImage<'a>>,
  16. }
  17. impl<'a> PHash<'a> {
  18. pub fn new(path: &'a Path, precision: &Precision, cache: &Option<Cache>) -> Self {
  19. PHash {
  20. prepared_image: Box::new(prepare_image(&path, &HashType::PHash, &precision, cache)),
  21. }
  22. }
  23. }
  24. impl<'a> PerceptualHash for PHash<'a> {
  25. /**
  26. * Calculate the phash of the provided prepared image
  27. *
  28. * # Return
  29. *
  30. * Returns a u64 representing the value of the hash
  31. */
  32. fn get_hash(&self, cache: &Option<Cache>) -> u64 {
  33. match self.prepared_image.image {
  34. Some(ref image) => {
  35. // Get the image data into a vector to perform the DFT on.
  36. let width = image.width() as usize;
  37. let height = image.height() as usize;
  38. // Get 2d data to 2d FFT/DFT
  39. // Either from the cache or calculate it
  40. // Pretty fast already, so caching doesn't make a huge difference
  41. // Atleast compared to opening and processing the images
  42. let data_matrix: Vec<Vec<f64>> = match *cache {
  43. Some(ref c) => {
  44. match c.get_matrix_from_cache(&Path::new(self.prepared_image.orig_path),
  45. width as u32) {
  46. Some(matrix) => matrix,
  47. None => {
  48. let matrix = create_data_matrix(width, height, &image);
  49. match c.put_matrix_in_cache(&Path::new(self.prepared_image.orig_path),
  50. width as u32,
  51. &matrix) {
  52. Ok(_) => {}
  53. Err(e) => println!("Unable to store matrix in cache. {}", e),
  54. };
  55. matrix
  56. }
  57. }
  58. }
  59. None => create_data_matrix(width, height, &image),
  60. };
  61. // Only need the top left quadrant
  62. let target_width = (width / 4) as usize;
  63. let target_height = (height / 4) as usize;
  64. let dft_width = (width / 4) as f64;
  65. let dft_height = (height / 4) as f64;
  66. // Calculate the mean
  67. let mut total = 0f64;
  68. for x in 0..target_width {
  69. for y in 0..target_height {
  70. total += data_matrix[x][y];
  71. }
  72. }
  73. let mean = total / (dft_width * dft_height);
  74. // Calculating a hash based on the mean
  75. let mut hash = 0u64;
  76. for x in 0..target_width {
  77. for y in 0..target_height {
  78. if data_matrix[x][y] >= mean {
  79. hash |= 1;
  80. } else {
  81. hash |= 0;
  82. }
  83. hash <<= 1;
  84. }
  85. }
  86. hash
  87. }
  88. None => 0u64,
  89. }
  90. }
  91. }
  92. fn create_data_matrix(width: usize,
  93. height: usize,
  94. image: &DynamicImage)
  95. -> Vec<Vec<f64>> {
  96. let mut data_matrix: Vec<Vec<f64>> = Vec::new();
  97. // Preparing the results
  98. for x in 0..width {
  99. data_matrix.push(Vec::new());
  100. for y in 0..height {
  101. let pos_x = x as u32;
  102. let pos_y = y as u32;
  103. data_matrix[x].push(image.get_pixel(pos_x, pos_y).channels()[0] as f64);
  104. }
  105. }
  106. // Perform the 2D DFT operation on our matrix
  107. calculate_2d_dft(&mut data_matrix);
  108. data_matrix
  109. }
  110. // Use a 1D DFT to cacluate the 2D DFT.
  111. //
  112. // This is achieved by calculating the DFT for each row, then calculating the
  113. // DFT for each column of DFT row data. This means that a 32x32 image with have
  114. // 1024 1D DFT operations performed on it. (Slightly caclulation intensive)
  115. //
  116. // This operation is in place on the data in the provided vector
  117. //
  118. // Inspired by:
  119. // http://www.inf.ufsc.br/~visao/khoros/html-dip/c5/s2/front-page.html
  120. //
  121. // Checked with:
  122. // http://calculator.vhex.net/post/calculator-result/2d-discrete-fourier-transform
  123. //
  124. fn calculate_2d_dft(data_matrix: &mut Vec<Vec<f64>>) {
  125. // println!("{:?}", data_matrix);
  126. let width = data_matrix.len();
  127. let height = data_matrix[0].len();
  128. let mut complex_data_matrix = Vec::with_capacity(width);
  129. // Perform DCT on the columns of data
  130. for x in 0..width {
  131. let mut column: Vec<f64> = Vec::with_capacity(height);
  132. for y in 0..height {
  133. column.push(data_matrix[x][y]);
  134. }
  135. // Perform the DCT on this column
  136. // println!("column[{}] before: {:?}", x, column);
  137. let forward_plan = dft::Plan::new(dft::Operation::Forward, column.len());
  138. column.transform(&forward_plan);
  139. let complex_column = dft::unpack(&column);
  140. // println!("column[{}] after: {:?}", x, complex_column);
  141. complex_data_matrix.push(complex_column);
  142. }
  143. // Perform DCT on the rows of data
  144. for y in 0..height {
  145. let mut row = Vec::with_capacity(width);
  146. for x in 0..width {
  147. row.push(complex_data_matrix[x][y]);
  148. }
  149. // Perform DCT on the row
  150. // println!("row[{}] before: {:?}", y, row);
  151. let forward_plan = dft::Plan::new(dft::Operation::Forward, row.len());
  152. row.transform(&forward_plan);
  153. // println!("row[{}] after: {:?}", y, row);
  154. // Put the row values back
  155. for x in 0..width {
  156. data_matrix[x][y] = round_float(row[x].re);
  157. }
  158. }
  159. }
  160. fn round_float(f: f64) -> f64 {
  161. if f >= super::FLOAT_PRECISION_MAX_1 || f <= super::FLOAT_PRECISION_MIN_1 {
  162. f
  163. } else if f >= super::FLOAT_PRECISION_MAX_2 || f <= super::FLOAT_PRECISION_MIN_2 {
  164. (f * 10_f64).round() / 10_f64
  165. } else if f >= super::FLOAT_PRECISION_MAX_3 || f <= super::FLOAT_PRECISION_MIN_3 {
  166. (f * 100_f64).round() / 100_f64
  167. } else if f >= super::FLOAT_PRECISION_MAX_4 || f <= super::FLOAT_PRECISION_MIN_4 {
  168. (f * 1000_f64).round() / 1000_f64
  169. } else if f >= super::FLOAT_PRECISION_MAX_5 || f <= super::FLOAT_PRECISION_MIN_5 {
  170. (f * 10000_f64).round() / 10000_f64
  171. } else {
  172. (f * 100000_f64).round() / 100000_f64
  173. }
  174. }
  175. #[test]
  176. fn test_2d_dft() {
  177. let mut test_matrix: Vec<Vec<f64>> = Vec::new();
  178. test_matrix.push(vec![1f64, 1f64, 1f64, 3f64]);
  179. test_matrix.push(vec![1f64, 2f64, 2f64, 1f64]);
  180. test_matrix.push(vec![1f64, 2f64, 2f64, 1f64]);
  181. test_matrix.push(vec![3f64, 1f64, 1f64, 1f64]);
  182. println!("{:?}", test_matrix[0]);
  183. println!("{:?}", test_matrix[1]);
  184. println!("{:?}", test_matrix[2]);
  185. println!("{:?}", test_matrix[3]);
  186. println!("Performing 2d DFT");
  187. calculate_2d_dft(&mut test_matrix);
  188. println!("{:?}", test_matrix[0]);
  189. println!("{:?}", test_matrix[1]);
  190. println!("{:?}", test_matrix[2]);
  191. println!("{:?}", test_matrix[3]);
  192. assert!(test_matrix[0][0] == 24_f64);
  193. assert!(test_matrix[0][1] == 0_f64);
  194. assert!(test_matrix[0][2] == 0_f64);
  195. assert!(test_matrix[0][3] == 0_f64);
  196. assert!(test_matrix[1][0] == 0_f64);
  197. assert!(test_matrix[1][1] == 0_f64);
  198. assert!(test_matrix[1][2] == -2_f64);
  199. assert!(test_matrix[1][3] == 2_f64);
  200. assert!(test_matrix[2][0] == 0_f64);
  201. assert!(test_matrix[2][1] == -2_f64);
  202. assert!(test_matrix[2][2] == -4_f64);
  203. assert!(test_matrix[2][3] == -2_f64);
  204. assert!(test_matrix[3][0] == 0_f64);
  205. assert!(test_matrix[3][1] == 2_f64);
  206. assert!(test_matrix[3][2] == -2_f64);
  207. assert!(test_matrix[3][3] == 0_f64);
  208. }