106 lines
4.6 KiB
Rust
106 lines
4.6 KiB
Rust
use autoquant::packing::ErrorFunction;
|
|
use graphene_core::quantization::*;
|
|
use graphene_core::raster::{Color, ImageFrame};
|
|
use graphene_core::Node;
|
|
|
|
/// The `GenerateQuantizationNode` encodes the brightness of each channel of the image as an integer number
|
|
/// signified by the samples parameter. This node is used to asses the loss of visual information when
|
|
/// quantizing the image using different fit functions.
|
|
pub struct GenerateQuantizationNode<N, M> {
|
|
samples: N,
|
|
function: M,
|
|
}
|
|
|
|
#[node_macro::node_fn(GenerateQuantizationNode)]
|
|
fn generate_quantization_fn(image_frame: ImageFrame<Color>, samples: u32, function: u32) -> [Quantization; 4] {
|
|
generate_quantization_from_image_frame(&image_frame)
|
|
}
|
|
|
|
pub fn generate_quantization_from_image_frame(image_frame: &ImageFrame<Color>) -> [Quantization; 4] {
|
|
let image = &image_frame.image;
|
|
|
|
let len = image.data.len().min(10000);
|
|
let data = image
|
|
.data
|
|
.iter()
|
|
.enumerate()
|
|
.filter(|(i, _)| i % (image.data.len() / len) == 0)
|
|
.flat_map(|(_, x)| vec![x.r() as f64, x.g() as f64, x.b() as f64, x.a() as f64])
|
|
.collect::<Vec<_>>();
|
|
generate_quantization(data, len)
|
|
}
|
|
fn generate_quantization(data: Vec<f64>, samples: usize) -> [Quantization; 4] {
|
|
let red = create_distribution(data.clone(), samples, 0);
|
|
let green = create_distribution(data.clone(), samples, 1);
|
|
let blue = create_distribution(data.clone(), samples, 2);
|
|
let alpha = create_distribution(data, samples, 3);
|
|
|
|
let fit_red = autoquant::calculate_error_function(&red, 1, &red);
|
|
let fit_green = autoquant::calculate_error_function(&green, 1, &green);
|
|
let fit_blue = autoquant::calculate_error_function(&blue, 1, &blue);
|
|
let fit_alpha = autoquant::calculate_error_function(&alpha, 1, &alpha);
|
|
let red_error: ErrorFunction<10> = autoquant::packing::ErrorFunction::new(fit_red.as_slice());
|
|
let green_error: ErrorFunction<10> = autoquant::packing::ErrorFunction::new(fit_green.as_slice());
|
|
let blue_error: ErrorFunction<10> = autoquant::packing::ErrorFunction::new(fit_blue.as_slice());
|
|
let alpha_error: ErrorFunction<10> = autoquant::packing::ErrorFunction::new(fit_alpha.as_slice());
|
|
let merged: ErrorFunction<20> = autoquant::packing::merge_error_functions(&red_error, &green_error);
|
|
let merged: ErrorFunction<30> = autoquant::packing::merge_error_functions(&merged, &blue_error);
|
|
let merged: ErrorFunction<40> = autoquant::packing::merge_error_functions(&merged, &alpha_error);
|
|
|
|
let bin_size = 8;
|
|
let mut distributions = [red, green, blue, alpha].into_iter();
|
|
|
|
let bits = &merged.bits[bin_size];
|
|
|
|
core::array::from_fn(|i| {
|
|
let fit = autoquant::models::OptimizedLin::new(distributions.next().unwrap(), (1 << bits[i]) - 1);
|
|
let parameters = fit.parameters();
|
|
Quantization::new(parameters[0] as f32, parameters[1] as f32, bits[i] as u32)
|
|
})
|
|
}
|
|
|
|
/*
|
|
// TODO: make this work with generic size parameters
|
|
fn generate_quantization<const N: usize>(data: Vec<f64>, samples: usize, channels: usize) -> [Quantization; N] {
|
|
let mut quantizations = Vec::new();
|
|
let mut merged_error: Option<ErrorFunction<10>> = None;
|
|
let bin_size = 32;
|
|
|
|
for i in 0..channels {
|
|
let channel_data = create_distribution(data.clone(), samples, i);
|
|
|
|
let fit = autoquant::calculate_error_function(&channel_data, 0, &channel_data);
|
|
let error: ErrorFunction<10> = autoquant::packing::ErrorFunction::new(fit.as_slice());
|
|
|
|
// Merge current error function with previous ones
|
|
merged_error = match merged_error {
|
|
Some(prev_error) => Some(autoquant::packing::merge_error_functions(&prev_error, &error)),
|
|
None => Some(error.clone()),
|
|
};
|
|
|
|
println!("Merged: {merged_error:?}");
|
|
|
|
let bits = merged_error.as_ref().unwrap().bits.iter().map(|x| x[i]).collect::<Vec<_>>();
|
|
let model_fit = autoquant::models::OptimizedLin::new(channel_data, 1 << bits[bin_size]);
|
|
let parameters = model_fit.parameters();
|
|
let quantization = Quantization::new(parameters[0] as f32, parameters[1] as u32, bits[bin_size] as u32);
|
|
|
|
quantizations.push(quantization);
|
|
}
|
|
|
|
core::array::from_fn(|x| quantizations[x])
|
|
}*/
|
|
|
|
fn create_distribution(data: Vec<f64>, samples: usize, channel: usize) -> Vec<(f64, f64)> {
|
|
let data: Vec<f64> = data.chunks(4 * (data.len() / (4 * samples.min(data.len() / 4)))).map(|x| x[channel]).collect();
|
|
let max = *data.iter().max_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal)).unwrap();
|
|
let data: Vec<f64> = data.iter().map(|x| x / max).collect();
|
|
dbg!(max);
|
|
// let data = autoquant::generate_normal_distribution(3.0, 1.1, 1000);
|
|
// data.iter_mut().for_each(|x| *x = x.abs());
|
|
let mut dist = autoquant::integrate_distribution(data);
|
|
autoquant::drop_duplicates(&mut dist);
|
|
let dist = autoquant::normalize_distribution(dist.as_slice());
|
|
dist
|
|
}
|