diff --git a/src/cv_compat.rs b/src/cv_compat.rs index c69f44b..7aac388 100644 --- a/src/cv_compat.rs +++ b/src/cv_compat.rs @@ -4,8 +4,7 @@ use ndarray::{Array3, ArrayView3}; /// OpenCV-compatible image decoding functionality pub mod imdecode { use super::*; - use image::ImageFormat; - use std::io::Cursor; + use image::DynamicImage; /// Image read flags matching OpenCV constants #[derive(Debug, Clone, Copy)] @@ -21,47 +20,68 @@ pub mod imdecode { /// Decode image from byte buffer (equivalent to cv2.imdecode) pub fn imdecode(buf: &[u8], flags: ImreadFlags) -> Result> { - let cursor = Cursor::new(buf); - let img = image::load( - cursor, - ImageFormat::from_extension("").unwrap_or(ImageFormat::Png), - ) - .map_err(|e| anyhow::anyhow!("Failed to decode image: {}", e))?; + let img = image::load_from_memory(buf) + .map_err(|e| anyhow::anyhow!("Failed to decode image: {}", e))?; match flags { ImreadFlags::ImreadGrayscale => { let gray_img = img.to_luma8(); - let (width, height) = gray_img.dimensions(); - - // Convert to 3-channel grayscale (RGB format) - let mut rgb_data = Array3::::zeros((height as usize, width as usize, 3)); - for y in 0..height { - for x in 0..width { - let pixel = gray_img.get_pixel(x, y); - let gray_val = pixel[0]; - rgb_data[[y as usize, x as usize, 0]] = gray_val; - rgb_data[[y as usize, x as usize, 1]] = gray_val; - rgb_data[[y as usize, x as usize, 2]] = gray_val; - } - } - Ok(rgb_data) + luma8_to_rgb_ndarray(gray_img) } - ImreadFlags::ImreadColor | ImreadFlags::ImreadUnchanged => { + ImreadFlags::ImreadColor => { let rgb_img = img.to_rgb8(); let (width, height) = rgb_img.dimensions(); + Array3::from_shape_vec((height as usize, width as usize, 3), rgb_img.into_raw()) + .map_err(|e| anyhow::anyhow!("Failed to reshape RGB image: {}", e)) + } + ImreadFlags::ImreadUnchanged => dynamic_image_to_ndarray(img), + } + } - let mut rgb_data = Array3::::zeros((height as usize, width as usize, 3)); - for y in 0..height { - for x in 0..width { - let pixel = rgb_img.get_pixel(x, y); - rgb_data[[y as usize, x as usize, 0]] = pixel[0]; - rgb_data[[y as usize, x as usize, 1]] = pixel[1]; - rgb_data[[y as usize, x as usize, 2]] = pixel[2]; - } - } - Ok(rgb_data) + fn dynamic_image_to_ndarray(img: DynamicImage) -> Result> { + match img.color().channel_count() { + 1 => { + let gray_img = img.to_luma8(); + let (width, height) = gray_img.dimensions(); + Array3::from_shape_vec((height as usize, width as usize, 1), gray_img.into_raw()) + .map_err(|e| anyhow::anyhow!("Failed to reshape grayscale image: {}", e)) + } + 2 => { + let gray_alpha_img = img.to_luma_alpha8(); + let (width, height) = gray_alpha_img.dimensions(); + Array3::from_shape_vec( + (height as usize, width as usize, 2), + gray_alpha_img.into_raw(), + ) + .map_err(|e| anyhow::anyhow!("Failed to reshape grayscale+alpha image: {}", e)) + } + 3 => { + let rgb_img = img.to_rgb8(); + let (width, height) = rgb_img.dimensions(); + Array3::from_shape_vec((height as usize, width as usize, 3), rgb_img.into_raw()) + .map_err(|e| anyhow::anyhow!("Failed to reshape RGB image: {}", e)) } + 4 => { + let rgba_img = img.to_rgba8(); + let (width, height) = rgba_img.dimensions(); + Array3::from_shape_vec((height as usize, width as usize, 4), rgba_img.into_raw()) + .map_err(|e| anyhow::anyhow!("Failed to reshape RGBA image: {}", e)) + } + channels => anyhow::bail!("Unsupported decoded image channel count: {}", channels), + } + } + + fn luma8_to_rgb_ndarray(gray_img: image::GrayImage) -> Result> { + let (width, height) = gray_img.dimensions(); + let gray_raw = gray_img.into_raw(); + let mut rgb_data = Vec::with_capacity(gray_raw.len() * 3); + + for gray in gray_raw { + rgb_data.extend_from_slice(&[gray, gray, gray]); } + + Array3::from_shape_vec((height as usize, width as usize, 3), rgb_data) + .map_err(|e| anyhow::anyhow!("Failed to reshape grayscale image: {}", e)) } } diff --git a/src/tests.rs b/src/tests.rs index 4e29c1f..1529269 100644 --- a/src/tests.rs +++ b/src/tests.rs @@ -332,8 +332,95 @@ mod edge_case_tests { #[cfg(test)] mod cv_compat_tests { - use crate::cv_compat::{cvt_color, ColorConversionCode}; + use crate::cv_compat::{cvt_color, imdecode, ColorConversionCode, ImreadFlags}; + use image::{DynamicImage, ImageBuffer, ImageFormat, Luma, LumaA, Rgb, Rgba}; use ndarray::Array3; + use std::io::Cursor; + + #[test] + fn test_imdecode_detects_jpeg_without_png_hint() { + let mut encoded = Cursor::new(Vec::new()); + let rgb = ImageBuffer::from_fn(3, 2, |x, y| Rgb([(x * 40) as u8, (y * 80) as u8, 200])); + + DynamicImage::ImageRgb8(rgb) + .write_to(&mut encoded, ImageFormat::Jpeg) + .unwrap(); + + let decoded = imdecode(encoded.get_ref(), ImreadFlags::ImreadColor).unwrap(); + assert_eq!(decoded.dim(), (2, 3, 3)); + } + + #[test] + fn test_imdecode_detects_webp_without_png_hint() { + let mut encoded = Cursor::new(Vec::new()); + let rgb = ImageBuffer::from_fn(4, 3, |x, y| Rgb([(x * 20) as u8, 100, (y * 30) as u8])); + + DynamicImage::ImageRgb8(rgb) + .write_to(&mut encoded, ImageFormat::WebP) + .unwrap(); + + let decoded = imdecode(encoded.get_ref(), ImreadFlags::ImreadColor).unwrap(); + assert_eq!(decoded.dim(), (3, 4, 3)); + } + + #[test] + fn test_imdecode_unchanged_preserves_alpha() { + let mut encoded = Cursor::new(Vec::new()); + let rgba = ImageBuffer::from_fn(2, 2, |x, y| { + let alpha = if (x + y) % 2 == 0 { 64 } else { 255 }; + Rgba([10 + x as u8, 20 + y as u8, 30, alpha]) + }); + + DynamicImage::ImageRgba8(rgba) + .write_to(&mut encoded, ImageFormat::Png) + .unwrap(); + + let decoded = imdecode(encoded.get_ref(), ImreadFlags::ImreadUnchanged).unwrap(); + assert_eq!(decoded.dim(), (2, 2, 4)); + assert_eq!(decoded[[0, 0, 3]], 64); + assert_eq!(decoded[[0, 1, 3]], 255); + assert_eq!(decoded[[1, 0, 3]], 255); + assert_eq!(decoded[[1, 1, 3]], 64); + } + + #[test] + fn test_imdecode_grayscale_returns_compat_rgb_shape() { + let mut encoded = Cursor::new(Vec::new()); + let gray = ImageBuffer::from_fn(2, 2, |x, y| Luma([10 + (x + y * 2) as u8])); + + DynamicImage::ImageLuma8(gray) + .write_to(&mut encoded, ImageFormat::Png) + .unwrap(); + + let decoded = imdecode(encoded.get_ref(), ImreadFlags::ImreadGrayscale).unwrap(); + assert_eq!(decoded.dim(), (2, 2, 3)); + assert_eq!(decoded[[0, 0, 0]], decoded[[0, 0, 1]]); + assert_eq!(decoded[[0, 0, 1]], decoded[[0, 0, 2]]); + assert_eq!(decoded[[1, 1, 0]], 13); + } + + #[test] + fn test_imdecode_unchanged_preserves_luma_alpha_channels() { + let mut encoded = Cursor::new(Vec::new()); + let gray_alpha = ImageBuffer::from_fn(2, 2, |x, y| { + let luma = 20 + (x + y * 2) as u8; + let alpha = if (x + y) % 2 == 0 { 32 } else { 200 }; + LumaA([luma, alpha]) + }); + + DynamicImage::ImageLumaA8(gray_alpha) + .write_to(&mut encoded, ImageFormat::Png) + .unwrap(); + + let decoded = imdecode(encoded.get_ref(), ImreadFlags::ImreadUnchanged).unwrap(); + assert_eq!(decoded.dim(), (2, 2, 2)); + assert_eq!(decoded[[0, 0, 0]], 20); + assert_eq!(decoded[[0, 0, 1]], 32); + assert_eq!(decoded[[0, 1, 0]], 21); + assert_eq!(decoded[[0, 1, 1]], 200); + assert_eq!(decoded[[1, 1, 0]], 23); + assert_eq!(decoded[[1, 1, 1]], 32); + } #[test] fn test_rgb_to_hsv_wraps_negative_red_sector_hues() {