Files
image-similarity/src/descriptors/dct.rs
T
2024-05-14 14:37:33 +02:00

202 lines
6.7 KiB
Rust

use image::{save_buffer, GenericImageView};
use std::f64::consts::{PI, SQRT_2};
use crate::descriptors::{Descriptor, DCT, print_matrix};
use log::{debug, log_enabled, Level};
impl DCT {
/// Returns a new DCT instance with a base quality DCT matrix.
pub fn new() -> DCT {
let quantization_matrix: [u8; 64] = [
16, 11, 10, 16, 24, 40, 51, 61,
12, 12, 14, 19, 26, 58, 60, 55,
14, 13, 16, 24, 40, 57, 69, 56,
14, 17, 22, 29, 51, 87, 80, 62,
18, 22, 37, 56, 6, 10, 103, 77,
24, 35, 55, 64, 8, 10, 113, 92,
49, 64, 78, 8, 10, 12, 12, 101,
72, 92, 95, 9, 11, 10, 103, 99,
];
DCT { quantization_matrix }
}
/// Builds DCT with given quality value
/// quality
pub fn with_quality(mut self, quality: u8) -> Self {
let mut quantization_matrix: [u8; 64] = [0; 64];
let scalar: f32 = match quality {
1..=49 => 5000.0/quality as f32,
50..=100 => 200.0 - 2.0*quality as f32,
_ => 100.0 // Invalid input: set to base quality
};
for i in 0..64 {
quantization_matrix[i] = ((scalar * self.quantization_matrix[i] as f32 + 50.0) / 100.0).floor() as u8;
if quantization_matrix[i] == 0 {
quantization_matrix[i] = 1;
}
}
self.quantization_matrix = quantization_matrix;
self
}
///
fn dct(&self, img: &image::DynamicImage) -> [f64; 64] {
let mut dct_values: [f64; 64] = [0.0; 64];
for u in 0..8 {
for v in 0..8 {
let k = (v*8)+u;
let mut alpha = 0.25;
if u == 0 {
alpha = alpha / SQRT_2
}
if v == 0 {
alpha = alpha / SQRT_2
}
let v: f64 = v as f64;
let u: f64 = u as f64;
let mut sum: f64 = 0.0;
for (x, y, pix) in img.pixels() {
let x: f64 = 1.0 + 2.0 * x as f64;
let y: f64 = 1.0 + 2.0 * y as f64;
let pixel = (pix[0] as i16 - 127) as f64;
sum +=
pixel *
(x*u*PI/16.0).cos() *
(y*v*PI/16.0).cos()
}
dct_values[k] = alpha * sum;
}
}
dct_values
}
fn idct(&self, dct_values: [f64; 64]) -> [u8; 64] {
let mut reconstructed: [u8; 64] = [0; 64];
for k in 0..64 {
let x = (k%8) as f64;
let y = (k/8) as f64;
let mut sum = 0.0;
for u in 0..8 {
for v in 0..8 {
let mut alpha = 1.0;
if u == 0 {
alpha = alpha / SQRT_2
}
if v == 0 {
alpha = alpha / SQRT_2
}
let uv = (v*8)+u;
let v = v as f64;
let u = u as f64;
sum +=
alpha *
dct_values[uv] *
((2.0 * x + 1.0) * u * PI / 16.0).cos() *
((2.0 * y + 1.0) * v * PI / 16.0).cos();
}
}
sum = 127.0 + (0.25 * sum).round();
reconstructed[k] = std::cmp::min(255_u8, sum as u8);
}
reconstructed
}
}
impl Descriptor for DCT {
fn info(&self) -> String {
"DCT".to_string()
}
fn describe(&self, img: &image::DynamicImage) -> u64 {
let resized = self.resize(img);
let mut dct_values = self.dct(&resized);
debug!("DCT-coefficients:\n {}", print_matrix(dct_values));
// Quantization:
//debug!("Using quantization matrix:\n{}", print_matrix(self.quantization_matrix));
for i in 0..64 {
dct_values[i] = (dct_values[i] / self.quantization_matrix[i] as f64).round();
}
debug!("DCT-coefficients, quantized:\n {}", print_matrix(dct_values));
if log_enabled!(Level::Debug) {
resized.save("resize.png").expect("Error saving file");
// De-quantization:
for i in 0..64 {
dct_values[i] = dct_values[i] * self.quantization_matrix[i] as f64;
}
//debug!("DCT-coefficients, de-quantized:\n {}", print_matrix(dct_values));
// Reconstruction original pixel values:
let reconstructed = self.idct(dct_values);
save_buffer(
"reconstructed.png",
&reconstructed,
8,
8,
image::ColorType::L8
).expect("Error saving buffer");
}
// Calculating descriptor from dct values:
// Zigzag order for our flattened array,
// first 28 elements only
let zigzag: [usize; 28] = [
0,
1, 8,
16, 9, 2,
3, 10, 17, 24,
32, 25, 18, 11, 4,
5, 12, 19, 26, 33, 40,
48, 41, 34, 27, 20, 13, 6,
];
// Mask that indicates if a dct coefficient is positive or negative
// By convention, when sign bit is 1, number is negative
let mut sign_mask: u64 = 0;
// If first horizontal AC coefficient is negative
// This might account for horizontal flips when applied to all horizontal coefficients.
let sign_mult = dct_values[1].signum();
// Mask that indicates if a coefficient is bigger or smaller than previous in order
let mut pearson_mask: u64 = 0;
let mut prev = dct_values[0];
for i in zigzag {
let cur = dct_values[i];
if cur > prev {
pearson_mask += 1;
}
prev = cur;
let signum = dct_values[i].signum();
// Only multiply sign if dct-coefficient contains a horizontal component.
if
//i % 8 != 0 && sign_mult * signum < 0.0
// ||
signum < 0.0
{
sign_mask += 1;
}
// Shift masks
sign_mask = sign_mask << 1;
pearson_mask = pearson_mask << 1;
}
debug!("Sign mask: {:028b}", sign_mask);
debug!("Pearson mask: {:028b}", pearson_mask);
let mut mask = sign_mask;
debug!("Mask: {:064b}", mask);
mask = mask << 28;
debug!("Mask: {:064b}", mask);
mask += pearson_mask;
debug!("Mask: {:064b}", mask);
mask = mask << 8;
debug!("Mask: {:064b}", mask);
// TODO: Do something with these last 8 bits.
mask
}
}