From 7eb42a27eeb425fd49268dde1f7165c8c6da5b94 Mon Sep 17 00:00:00 2001 From: colbyn Date: Sun, 8 Mar 2020 21:39:47 -0600 Subject: [PATCH] Building out prerequisites for a new idea --- classifier/Cargo.toml | 4 + classifier/src/classify.rs | 113 +++++++++++++++++++++++++++ classifier/src/codec/jpeg.rs | 127 ++++++++++++++++++++++++++++++ classifier/src/codec/mod.rs | 1 + classifier/src/color/format.rs | 137 +++++++++++++++++++++++++++++++++ classifier/src/color/mod.rs | 3 +- classifier/src/main.rs | 5 +- 7 files changed, 388 insertions(+), 2 deletions(-) create mode 100644 classifier/src/classify.rs create mode 100644 classifier/src/codec/jpeg.rs create mode 100644 classifier/src/codec/mod.rs create mode 100644 classifier/src/color/format.rs diff --git a/classifier/Cargo.toml b/classifier/Cargo.toml index ebd7364..c86c761 100644 --- a/classifier/Cargo.toml +++ b/classifier/Cargo.toml @@ -18,3 +18,7 @@ serde_json = "^1.0" exoquant = "0.2.0" lodepng = "2.5.0" hsl-ish = "0.1.0" +av-metrics = { version = "0.4", features = ["serde"] } +y4m = "0.5" +dcv-color-primitives = "0.1.4" +lazy_static = "1.4.0" diff --git a/classifier/src/classify.rs b/classifier/src/classify.rs new file mode 100644 index 0000000..e998cdc --- /dev/null +++ b/classifier/src/classify.rs @@ -0,0 +1,113 @@ +use std::io::Cursor; +use av_metrics::video::*; +use image::{DynamicImage, GenericImage, GenericImageView}; +// use clap::{App, Arg}; +// use console::style; +// use maplit::hashmap; +// use serde::Serialize; +// use std::collections::HashMap; +// use std::error::Error; +// use std::fs::File; +// use std::path::Path; +// use std::process::exit; + + + +pub fn run() { + let input_path = "assets/samples/pexels-photo-1153655.jpeg"; + let quality = 10; + + // SOURCE-1 + let source1 = ::image::open(input_path).expect("source image"); + + // SOURCE-2 + let source2 = ::image::open(input_path).expect("source image"); + let source2 = unsafe {crate::codec::jpeg::encode(&source2, quality)}; + std::fs::write("test.jpeg", &source2); + let source2 = ::image::load_from_memory_with_format(&source2, ::image::ImageFormat::Jpeg).expect("decode jpeg buffer"); + + // RESIZE + let source1 = source1.resize(300, 300, ::image::imageops::FilterType::Lanczos3); + let source2 = source2.resize(300, 300, ::image::imageops::FilterType::Lanczos3); + + // START METRICS + println!("computing metrics: {}", quality); + + // RUN SSIM + let to_encoded = |source: &DynamicImage| { + let encoded = { + let ([y, u, v], width, height) = crate::color::format::to_yuv420p(&source); + let y4m_frame = y4m::Frame::new([&y, &u, &v], None); + let encoder = y4m::encode( + width, + height, + y4m::Ratio::new(0, 1), + ); + let mut buffer = Vec::::new(); + { + let encoder = encoder.with_colorspace(y4m::Colorspace::C420); + let mut encoder = encoder.write_header(&mut buffer).expect("y4m encoder init"); + encoder.write_frame(&y4m_frame).expect("write y4m frame"); + } + buffer + }; + std::io::Cursor::new(encoded) + }; + + let ssim = { + let mut decoder1 = to_encoded(&source1); + let mut decoder2 = to_encoded(&source2); + + let mut decoder1: y4m::Decoder>> = y4m::Decoder::new(&mut decoder1).expect("init y4m decoder"); + let mut decoder2: y4m::Decoder>> = y4m::Decoder::new(&mut decoder2).expect("init y4m decoder"); + av_metrics::video::ssim::calculate_video_ssim( + &mut decoder1, + &mut decoder2, + None, + ).expect("calculate_video_ssim") + }; + + let msssim = { + let mut decoder1 = to_encoded(&source1); + let mut decoder2 = to_encoded(&source2); + + let mut decoder1: y4m::Decoder>> = y4m::Decoder::new(&mut decoder1).expect("init y4m decoder"); + let mut decoder2: y4m::Decoder>> = y4m::Decoder::new(&mut decoder2).expect("init y4m decoder"); + av_metrics::video::ssim::calculate_video_msssim( + &mut decoder1, + &mut decoder2, + None, + ).expect("calculate_video_msssim") + }; + + let ciede = { + let mut decoder1 = to_encoded(&source1); + let mut decoder2 = to_encoded(&source2); + + let mut decoder1: y4m::Decoder>> = y4m::Decoder::new(&mut decoder1).expect("init y4m decoder"); + let mut decoder2: y4m::Decoder>> = y4m::Decoder::new(&mut decoder2).expect("init y4m decoder"); + av_metrics::video::ciede::calculate_video_ciede( + &mut decoder1, + &mut decoder2, + None, + ).expect("calculate_video_ciede") + }; + + let psnr_hvs = { + let mut decoder1 = to_encoded(&source1); + let mut decoder2 = to_encoded(&source2); + + let mut decoder1: y4m::Decoder>> = y4m::Decoder::new(&mut decoder1).expect("init y4m decoder"); + let mut decoder2: y4m::Decoder>> = y4m::Decoder::new(&mut decoder2).expect("init y4m decoder"); + av_metrics::video::psnr_hvs::calculate_video_psnr_hvs( + &mut decoder1, + &mut decoder2, + None, + ).expect("calculate_video_psnr_hvs") + }; + + println!("ssim: {:#?}", ssim); + println!("msssim: {:#?}", msssim); + println!("ciede: {:#?}", ciede); + println!("psnr_hvs: {:#?}", psnr_hvs); +} diff --git a/classifier/src/codec/jpeg.rs b/classifier/src/codec/jpeg.rs new file mode 100644 index 0000000..6ff7043 --- /dev/null +++ b/classifier/src/codec/jpeg.rs @@ -0,0 +1,127 @@ +// This Source Code Form is subject to the terms of the Mozilla Public +// License, v. 2.0. If a copy of the MPL was not distributed with this +// file, You can obtain one at https://mozilla.org/MPL/2.0/. +use std::path::PathBuf; +use std::convert::From; +use std::ffi::{CStr, CString}; +use std::os::raw::{c_char, c_int}; +use libc::{size_t, c_float, c_void}; +use serde::{Serialize, Deserialize}; +use image::{ + GenericImage, + DynamicImage, + GenericImageView, + ColorType, + Pixel, +}; +use rayon::prelude::*; +use itertools::Itertools; + +/////////////////////////////////////////////////////////////////////////////// +// MOZJPEG FFI HELPERS +/////////////////////////////////////////////////////////////////////////////// + +#[allow(non_snake_case)] +const TRUE: mozjpeg_sys::boolean = true as mozjpeg_sys::boolean; +#[allow(non_snake_case)] +const FALSE: mozjpeg_sys::boolean = false as mozjpeg_sys::boolean; + +const COLOR_SPACE: mozjpeg_sys::J_COLOR_SPACE = mozjpeg_sys::J_COLOR_SPACE::JCS_RGB; +const COLOR_SPACE_COMPONENTS: libc::c_int = 3 as libc::c_int; + + +/////////////////////////////////////////////////////////////////////////////// +// MOZJPEG ENCODER +/////////////////////////////////////////////////////////////////////////////// + + +pub unsafe fn encode(source: &DynamicImage, quality: u8) -> Vec { + /////////////////////////////////////////////////////////////////////////// + // INPUT + /////////////////////////////////////////////////////////////////////////// + let rgb_source = source + .to_rgb() + .pixels() + .flat_map(|x| x.0.to_vec()) + .collect::>(); + let (width, height) = source.dimensions(); + + /////////////////////////////////////////////////////////////////////////// + // INIT ENCODER CONTEXT + /////////////////////////////////////////////////////////////////////////// + let mut err = std::mem::zeroed(); + let mut cinfo: mozjpeg_sys::jpeg_compress_struct = std::mem::zeroed(); + let mut outbuffer: *mut libc::c_uchar = std::ptr::null_mut(); + let mut outsize: libc::c_ulong = 0; + + cinfo.common.err = mozjpeg_sys::jpeg_std_error(&mut err); + mozjpeg_sys::jpeg_create_compress(&mut cinfo); + mozjpeg_sys::jpeg_mem_dest(&mut cinfo, &mut outbuffer, &mut outsize); + + /////////////////////////////////////////////////////////////////////////// + // ENCODER CONFIG + /////////////////////////////////////////////////////////////////////////// + cinfo.image_width = width; + cinfo.image_height = height; + cinfo.input_components = COLOR_SPACE_COMPONENTS; + let row_stride = cinfo.image_width as usize * cinfo.input_components as usize; + cinfo.in_color_space = COLOR_SPACE; + mozjpeg_sys::jpeg_set_defaults(&mut cinfo); + cinfo.dct_method = mozjpeg_sys::J_DCT_METHOD::JDCT_ISLOW; + cinfo.write_JFIF_header = FALSE; + cinfo.optimize_coding = TRUE; + mozjpeg_sys::jpeg_simple_progression(&mut cinfo); + mozjpeg_sys::jpeg_c_set_bool_param(&mut cinfo, mozjpeg_sys::JBOOLEAN_USE_SCANS_IN_TRELLIS, TRUE); + mozjpeg_sys::jpeg_c_set_bool_param(&mut cinfo, mozjpeg_sys::JBOOLEAN_USE_LAMBDA_WEIGHT_TBL, TRUE); + mozjpeg_sys::jpeg_set_quality(&mut cinfo, quality as i32, TRUE); + + /////////////////////////////////////////////////////////////////////////// + // GO! + /////////////////////////////////////////////////////////////////////////// + mozjpeg_sys::jpeg_start_compress(&mut cinfo, TRUE); + while cinfo.next_scanline < cinfo.image_height { + let offset = cinfo.next_scanline as usize * row_stride; + let jsamparray = [rgb_source[offset..].as_ptr()]; + mozjpeg_sys::jpeg_write_scanlines(&mut cinfo, jsamparray.as_ptr(), 1); + } + mozjpeg_sys::jpeg_finish_compress(&mut cinfo); + mozjpeg_sys::jpeg_destroy_compress(&mut cinfo); + + /////////////////////////////////////////////////////////////////////////// + // OUTPUT + /////////////////////////////////////////////////////////////////////////// + let output_data = std::slice::from_raw_parts(outbuffer, outsize as usize).to_vec(); + + /////////////////////////////////////////////////////////////////////////// + // CLEANUP + /////////////////////////////////////////////////////////////////////////// + if !outbuffer.is_null() { + // FREE MEMORY DEST + libc::free(outbuffer as *mut mozjpeg_sys::c_void); + outbuffer = std::ptr::null_mut(); + outsize = 0; + } + + /////////////////////////////////////////////////////////////////////////// + // DONE + /////////////////////////////////////////////////////////////////////////// + output_data +} + +/////////////////////////////////////////////////////////////////////////////// +// OPT +/////////////////////////////////////////////////////////////////////////////// + + + +/////////////////////////////////////////////////////////////////////////////// +// DEV +/////////////////////////////////////////////////////////////////////////////// + +// pub fn run() { +// let input_path = "assets/samples/ceiling.jpeg"; +// let source = ::image::open(input_path).expect("source image"); +// let (encoded, report) = OptContext::from_image(source).run_search(false); +// println!("results: {:#?}", report); +// std::fs::write("assets/output/test.jpeg", encoded); +// } \ No newline at end of file diff --git a/classifier/src/codec/mod.rs b/classifier/src/codec/mod.rs new file mode 100644 index 0000000..b3fe865 --- /dev/null +++ b/classifier/src/codec/mod.rs @@ -0,0 +1 @@ +pub mod jpeg; \ No newline at end of file diff --git a/classifier/src/color/format.rs b/classifier/src/color/format.rs new file mode 100644 index 0000000..c5ba13f --- /dev/null +++ b/classifier/src/color/format.rs @@ -0,0 +1,137 @@ +use std::sync::{Arc, Mutex}; +use image::{DynamicImage, GenericImage, GenericImageView}; +use dcv_color_primitives::{convert_image, ColorSpace, ImageFormat, PixelFormat}; + +fn ensure_even_reslution(source: &DynamicImage) -> DynamicImage { + let (width, height) = source.dimensions(); + // ENSURE EVEN + let even_width = (width % 2) == 0; + let even_height = (height % 2) == 0; + if (!even_width) || (!even_height) { + let new_width = { + if !even_width { + width - 1 + } else { + width + } + }; + let new_height = { + if !even_height { + height - 1 + } else { + height + } + }; + let new_image = source + .clone() + .crop(0, 0, new_width, new_height); + new_image + } else { + source.clone() + } +} + + +pub fn to_nv12(source: &DynamicImage) -> (Vec, usize, usize) { + // ENSURE VALID INPUT IMAGE + let source = ensure_even_reslution(source); + // SETUP + dcv_color_primitives::initialize(); + let (mut width, height) = source.dimensions(); + + // ALLOCATE INPUT + let source_buffer: Vec = { + source + .to_bgra() + .pixels() + .flat_map(|px: &::image::Bgra| vec![ + px.0[0], + px.0[1], + px.0[2], + px.0[3], + ]) + .collect::>() + }; + let input_data: &[&[u8]] = &[&source_buffer[..]]; + let src_format = ImageFormat { + pixel_format: PixelFormat::Bgra, + color_space: ColorSpace::Lrgb, + num_planes: 1, + }; + + // ALLOCATE OUTPUT + let dst_size: usize = 3 * (width as usize) * (height as usize) / 2; + let mut output_buffer = vec![0u8; dst_size]; + let output_data: &mut [&mut [u8]] = &mut [&mut output_buffer[..]]; + let dst_format = ImageFormat { + pixel_format: PixelFormat::Nv12, + color_space: ColorSpace::Bt601, + num_planes: 1, + }; + + // GO! + convert_image( + width, + height, + &src_format, + None, + input_data, + &dst_format, + None, + output_data, + ).expect("convert rgba source to nv12"); + + // DONE + assert!(output_data.len() == 1); + ( + output_data[0].to_owned(), + source.width() as usize, + source.height() as usize, + ) +} + +pub fn to_yuv420p(source: &DynamicImage) -> ([Vec; 3], usize, usize) { + use itertools::Itertools; + let (mut nv12, width, height) = to_nv12(source); + std::fs::write("test.nv12.yuv", &nv12); + let y_size: usize = { + (width * height) as usize + }; + let uv_size_interleaved: usize = { + (width * height / 2) as usize + }; + let uv_size_planar: usize = { + (width * height / 4) as usize + }; + // println!("nv12: {}", nv12.len()); + // println!("y_size: {}", y_size); + // println!("uv_size_interleaved: {}", uv_size_interleaved); + assert!(nv12.len() == y_size + uv_size_interleaved); + let y = nv12 + .drain(0 .. y_size) + .collect::>(); + assert!(nv12.len() == uv_size_interleaved); + let (u, v) = nv12 + .into_iter() + .chunks(2) + .into_iter() + .map(|uv| { + let uv = uv.collect::>(); + assert!(uv.len() == 2); + let u = uv[0]; + let v = uv[1]; + (u, v) + }) + .unzip::<_, _, Vec<_>, Vec<_>>(); + assert!(u.len() + v.len() == uv_size_interleaved); + assert!(u.len() == uv_size_planar); + assert!(v.len() == uv_size_planar); + ([y, u, v], width, height) +} + + +// pub fn to_y4m(source: &DynamicImage) -> y4m::Frame { +// let [y, u, v] = to_yuv420p(source); +// let frame = y4m::Frame::new([&y, &u, &v], None); +// frame +// } \ No newline at end of file diff --git a/classifier/src/color/mod.rs b/classifier/src/color/mod.rs index 62d7a5c..c3f6541 100644 --- a/classifier/src/color/mod.rs +++ b/classifier/src/color/mod.rs @@ -1,2 +1,3 @@ pub mod palette; -pub mod quant; \ No newline at end of file +pub mod quant; +pub mod format; \ No newline at end of file diff --git a/classifier/src/main.rs b/classifier/src/main.rs index 503adfa..6d9ceaf 100644 --- a/classifier/src/main.rs +++ b/classifier/src/main.rs @@ -1,7 +1,10 @@ #![allow(unused)] pub mod process; pub mod color; +pub mod classify; +pub mod codec; fn main() { - process::run(); + // process::run(); + classify::run(); }