From d7d782946e5c2bf8900eec5f56c890eb54b02b2b Mon Sep 17 00:00:00 2001 From: Derek Buitenhuis Date: Thu, 14 May 2020 16:44:56 +0100 Subject: [PATCH] capi: Expose frame buffer API Signed-off-by: Derek Buitenhuis --- Cargo.lock | 2 +- Cargo.toml | 2 +- av_metrics/Cargo.toml | 2 +- av_metrics/cbindgen.toml | 9 + av_metrics/src/capi.rs | 328 ++++++++++++++++++++++++++++++++++++ av_metrics/src/video/mod.rs | 1 + 6 files changed, 341 insertions(+), 3 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 7c52678..2e37ca9 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -40,7 +40,7 @@ dependencies = [ [[package]] name = "av-metrics" -version = "0.5.0" +version = "0.6.0" dependencies = [ "criterion 0.3.2 (registry+https://github.com/rust-lang/crates.io-index)", "itertools 0.9.0 (registry+https://github.com/rust-lang/crates.io-index)", diff --git a/Cargo.toml b/Cargo.toml index 8b46bf1..34aac47 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -9,4 +9,4 @@ members = [ opt-level = 1 [replace] -"av-metrics:0.5.0" = { path = "av_metrics" } +"av-metrics:0.6.0" = { path = "av_metrics" } diff --git a/av_metrics/Cargo.toml b/av_metrics/Cargo.toml index 87d8970..2b1df4b 100644 --- a/av_metrics/Cargo.toml +++ b/av_metrics/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "av-metrics" -version = "0.5.0" +version = "0.6.0" authors = ["Josh Holmer "] edition = "2018" description = "A collection of algorithms for measuring audio/video metrics" diff --git a/av_metrics/cbindgen.toml b/av_metrics/cbindgen.toml index 358aa7c..6280803 100644 --- a/av_metrics/cbindgen.toml +++ b/av_metrics/cbindgen.toml @@ -6,5 +6,14 @@ tab_width = 4 style = "Type" language = "C" +[parse] +parse_deps = true +include = ['av_metrics', 'v_frame'] + [export] prefix = "AVM" +item_types = ["enums", "structs", "unions", "typedefs", "opaque", "functions"] + +[enum] +rename_variants = "ScreamingSnakeCase" +prefix_with_name = true diff --git a/av_metrics/src/capi.rs b/av_metrics/src/capi.rs index c57198f..42a1670 100644 --- a/av_metrics/src/capi.rs +++ b/av_metrics/src/capi.rs @@ -3,14 +3,21 @@ extern crate libc; use libc::c_char; +use libc::ptrdiff_t; use std::ffi::CStr; use std::fs::File; use std::os::raw::c_int; use std::path::{Path, PathBuf}; use std::ptr::null; +use std::slice; +use crate::video as vid; use crate::video::*; +type ChromaSamplePosition = vid::ChromaSamplePosition; +type ChromaSampling = vid::ChromaSampling; +type Rational = vid::Rational; + #[derive(Debug, Clone, Copy)] enum InputType { Video(VideoContainer), @@ -436,6 +443,327 @@ pub unsafe extern fn avm_calculate_frame_ciede( value } +fn calculate_frame_buf_tmpl( + frame1: [&[u8]; 3], + frame1_strides: [ptrdiff_t; 3], + frame2: [&[u8]; 3], + frame2_strides: [ptrdiff_t; 3], + width: u32, + height: u32, + bitdepth: u8, + _chroma_pos: ChromaSamplePosition, + subsampling: ChromaSampling, + _pixel_aspect_ratio: Rational, + metric: &str, +) -> (*const Context, f64) { + let (xdec, ydec) = subsampling.get_decimation().unwrap_or((1, 1)); + let planes = if subsampling == ChromaSampling::Cs400 { + 1 + } else { + 3 + }; + let bw = if bitdepth == 8 { 1 } else { 2 }; + + let mut fi1 = FrameInfo { + planes: [ + Plane::::new(width as usize, height as usize, 0, 0, 0, 0), + Plane::::new( + (width as usize) >> xdec, + (height as usize) >> ydec, + xdec, + ydec, + 0, + 0, + ), + Plane::::new( + (width as usize) >> xdec, + (height as usize) >> ydec, + xdec, + ydec, + 0, + 0, + ), + ], + bit_depth: bitdepth as usize, + chroma_sampling: subsampling, + }; + let mut fi2 = FrameInfo { + planes: [ + Plane::::new(width as usize, height as usize, 0, 0, 0, 0), + Plane::::new( + (width as usize) >> xdec, + (height as usize) >> ydec, + xdec, + ydec, + 0, + 0, + ), + Plane::::new( + (width as usize) >> xdec, + (height as usize) >> ydec, + xdec, + ydec, + 0, + 0, + ), + ], + bit_depth: bitdepth as usize, + chroma_sampling: subsampling, + }; + + for p in 0..planes { + fi1.planes[p].copy_from_raw_u8(frame1[p], frame1_strides[p] as usize, bw); + fi2.planes[p].copy_from_raw_u8(frame2[p], frame2_strides[p] as usize, bw); + } + + if metric == "ciede" { + if let Ok(val) = ciede::calculate_frame_ciede(&fi1, &fi2) { + return (null(), val); + } + } + + let val = match metric { + "psnr" => psnr::calculate_frame_psnr(&fi1, &fi2), + "psnr_hvs" => psnr_hvs::calculate_frame_psnr_hvs(&fi1, &fi2), + "ssim" => ssim::calculate_frame_ssim(&fi1, &fi2), + "msssim" => ssim::calculate_frame_msssim(&fi1, &fi2), + _ => unimplemented!("unknown metric"), + }; + + if let Ok(metrics) = val { + let ctx = Context { + y: metrics.y, + u: metrics.u, + v: metrics.v, + avg: metrics.avg, + }; + let boxed = Box::new(ctx); + return (Box::into_raw(boxed), 0.0); + } + + (null(), -1.0) +} + +unsafe fn calculate_frame_buf_internal( + frame1: [*const u8; 3], + frame1_strides: [ptrdiff_t; 3], + frame2: [*const u8; 3], + frame2_strides: [ptrdiff_t; 3], + width: u32, + height: u32, + bitdepth: u8, + chroma_pos: ChromaSamplePosition, + subsampling: ChromaSampling, + pixel_aspect_ratio: Rational, + metric: &str, +) -> (*const Context, f64) { + let (_cw, ch) = subsampling.get_chroma_dimensions(width as usize, height as usize); + + let luma_len1 = (frame1_strides[0] as usize) * (height as usize); + let luma_slice1 = slice::from_raw_parts(frame1[0], luma_len1); + let chroma_u_len1 = (frame1_strides[1] as usize) * (ch as usize); + let chroma_u_slice1 = slice::from_raw_parts(frame1[1], chroma_u_len1); + let chroma_v_len1 = (frame1_strides[2] as usize) * (ch as usize); + let chroma_v_slice1 = slice::from_raw_parts(frame1[2], chroma_v_len1); + + let luma_len2 = (frame2_strides[0] as usize) * (height as usize); + let luma_slice2 = slice::from_raw_parts(frame2[0], luma_len2); + let chroma_u_len2 = (frame2_strides[1] as usize) * (ch as usize); + let chroma_u_slice2 = slice::from_raw_parts(frame2[1], chroma_u_len2); + let chroma_v_len2 = (frame2_strides[2] as usize) * (ch as usize); + let chroma_v_slice2 = slice::from_raw_parts(frame2[2], chroma_v_len2); + + if bitdepth == 8 { + calculate_frame_buf_tmpl::( + [&luma_slice1, &chroma_u_slice1, &chroma_v_slice1], + frame1_strides, + [&luma_slice2, &chroma_u_slice2, &chroma_v_slice2], + frame2_strides, + width, + height, + bitdepth, + chroma_pos, + subsampling, + pixel_aspect_ratio, + metric, + ) + } else { + calculate_frame_buf_tmpl::( + [&luma_slice1, &chroma_u_slice1, &chroma_v_slice1], + frame1_strides, + [&luma_slice2, &chroma_u_slice2, &chroma_v_slice2], + frame2_strides, + width, + height, + bitdepth, + chroma_pos, + subsampling, + pixel_aspect_ratio, + metric, + ) + } +} + +/// Calculate the `ciede` metric between two frame buffers +/// +/// Returns the correct `ciede` value or `-1` on errors +#[no_mangle] +pub unsafe extern fn avm_calculate_frame_buf_ciede( + frame1: [*const u8; 3], + frame1_strides: [ptrdiff_t; 3], + frame2: [*const u8; 3], + frame2_strides: [ptrdiff_t; 3], + width: u32, + height: u32, + bitdepth: u8, + chroma_pos: ChromaSamplePosition, + subsampling: ChromaSampling, + pixel_aspect_ratio: Rational, +) -> f64 { + let (_ctx, val) = calculate_frame_buf_internal( + frame1, + frame1_strides, + frame2, + frame2_strides, + width, + height, + bitdepth, + chroma_pos, + subsampling, + pixel_aspect_ratio, + "ciede", + ); + val +} + +/// Calculate the `ssim` metric between two frame buffers +/// +/// Returns the correct `ssim` value or `NULL` on errors +#[no_mangle] +pub unsafe extern fn avm_calculate_frame_buf_ssim( + frame1: [*const u8; 3], + frame1_strides: [ptrdiff_t; 3], + frame2: [*const u8; 3], + frame2_strides: [ptrdiff_t; 3], + width: u32, + height: u32, + bitdepth: u8, + chroma_pos: ChromaSamplePosition, + subsampling: ChromaSampling, + pixel_aspect_ratio: Rational, +) -> *const Context { + let (ctx, _val) = calculate_frame_buf_internal( + frame1, + frame1_strides, + frame2, + frame2_strides, + width, + height, + bitdepth, + chroma_pos, + subsampling, + pixel_aspect_ratio, + "ssim", + ); + ctx +} + +/// Calculate the `msssim` metric between two frame buffers +/// +/// Returns the correct `msssim` value or `NULL` on errors +#[no_mangle] +pub unsafe extern fn avm_calculate_frame_buf_msssim( + frame1: [*const u8; 3], + frame1_strides: [ptrdiff_t; 3], + frame2: [*const u8; 3], + frame2_strides: [ptrdiff_t; 3], + width: u32, + height: u32, + bitdepth: u8, + chroma_pos: ChromaSamplePosition, + subsampling: ChromaSampling, + pixel_aspect_ratio: Rational, +) -> *const Context { + let (ctx, _val) = calculate_frame_buf_internal( + frame1, + frame1_strides, + frame2, + frame2_strides, + width, + height, + bitdepth, + chroma_pos, + subsampling, + pixel_aspect_ratio, + "msssim", + ); + ctx +} + +/// Calculate the `psnr` metric between two frame buffers +/// +/// Returns the correct `psnr` value or `NULL` on errors +#[no_mangle] +pub unsafe extern fn avm_calculate_frame_buf_psnr( + frame1: [*const u8; 3], + frame1_strides: [ptrdiff_t; 3], + frame2: [*const u8; 3], + frame2_strides: [ptrdiff_t; 3], + width: u32, + height: u32, + bitdepth: u8, + chroma_pos: ChromaSamplePosition, + subsampling: ChromaSampling, + pixel_aspect_ratio: Rational, +) -> *const Context { + let (ctx, _val) = calculate_frame_buf_internal( + frame1, + frame1_strides, + frame2, + frame2_strides, + width, + height, + bitdepth, + chroma_pos, + subsampling, + pixel_aspect_ratio, + "psnr", + ); + ctx +} + +/// Calculate the `psnr_hvs` metric between two frame buffers +/// +/// Returns the correct `psnr_hvs` value or `NULL` on errors +#[no_mangle] +pub unsafe extern fn avm_calculate_frame_buf_psnr_hvs( + frame1: [*const u8; 3], + frame1_strides: [ptrdiff_t; 3], + frame2: [*const u8; 3], + frame2_strides: [ptrdiff_t; 3], + width: u32, + height: u32, + bitdepth: u8, + chroma_pos: ChromaSamplePosition, + subsampling: ChromaSampling, + pixel_aspect_ratio: Rational, +) -> *const Context { + let (ctx, _val) = calculate_frame_buf_internal( + frame1, + frame1_strides, + frame2, + frame2_strides, + width, + height, + bitdepth, + chroma_pos, + subsampling, + pixel_aspect_ratio, + "psnr_hvs", + ); + ctx +} + /// Drop the metric context /// /// This function drops the context and free the memory diff --git a/av_metrics/src/video/mod.rs b/av_metrics/src/video/mod.rs index b4f1197..bc8ae3d 100644 --- a/av_metrics/src/video/mod.rs +++ b/av_metrics/src/video/mod.rs @@ -94,6 +94,7 @@ impl ChromaWeight for ChromaSampling { /// Sample position for subsampled chroma #[derive(Copy, Clone, Debug, PartialEq)] +#[repr(C)] pub enum ChromaSamplePosition { /// The source video transfer function is not signaled. This crate will assume /// no transformation needs to be done on this data, but there is a risk of metric