use super::simd_accelerated;
use crate::*;
use scirs2_core::ndarray::{Array2, Array3, ArrayView2, ArrayView3};
pub fn extract_patches_2d(
image: &ArrayView2<Float>,
patch_size: (usize, usize),
max_patches: Option<usize>,
random_state: Option<u64>,
) -> SklResult<Array3<Float>> {
let (img_height, img_width) = image.dim();
let (patch_height, patch_width) = patch_size;
if patch_height == 0 || patch_width == 0 {
return Err(SklearsError::InvalidInput(
"Patch size dimensions must be positive".to_string(),
));
}
if patch_height > img_height || patch_width > img_width {
return Err(SklearsError::InvalidInput(
"Patch size cannot be larger than image size".to_string(),
));
}
let max_row = img_height - patch_height + 1;
let max_col = img_width - patch_width + 1;
let total_patches = max_row * max_col;
if total_patches == 0 {
return Ok(Array3::zeros((0, patch_height, patch_width)));
}
let n_patches = max_patches.unwrap_or(total_patches).min(total_patches);
let positions =
generate_patch_positions(max_row, max_col, n_patches, max_patches, random_state)?;
extract_patches_optimized(image, patch_size, &positions)
}
pub fn reconstruct_from_patches_2d(
patches: &ArrayView3<Float>,
image_size: (usize, usize),
) -> SklResult<Array2<Float>> {
let (height, width) = image_size;
if height == 0 || width == 0 {
return Err(SklearsError::InvalidInput(
"Image dimensions must be positive".to_string(),
));
}
simd_accelerated::simd_reconstruct_from_patches_2d(patches, image_size)
}
#[derive(Debug, Clone)]
pub struct PatchExtractor {
patch_size: (usize, usize),
max_patches: Option<usize>,
random_state: Option<u64>,
sampling_strategy: SamplingStrategy,
}
#[derive(Debug, Clone)]
pub enum SamplingStrategy {
UniformGrid,
RandomUniform,
Exhaustive,
Custom(Vec<(usize, usize)>),
}
impl PatchExtractor {
pub fn new() -> Self {
Self {
patch_size: (8, 8),
max_patches: None,
random_state: None,
sampling_strategy: SamplingStrategy::UniformGrid,
}
}
pub fn patch_size(mut self, patch_size: (usize, usize)) -> Self {
assert!(
patch_size.0 > 0 && patch_size.1 > 0,
"Patch dimensions must be positive"
);
self.patch_size = patch_size;
self
}
pub fn max_patches(mut self, max_patches: Option<usize>) -> Self {
self.max_patches = max_patches;
self
}
pub fn random_state(mut self, random_state: Option<u64>) -> Self {
self.random_state = random_state;
self
}
pub fn sampling_strategy(mut self, strategy: SamplingStrategy) -> Self {
self.sampling_strategy = strategy;
self
}
pub fn extract(&self, image: &ArrayView2<Float>) -> SklResult<Array3<Float>> {
extract_patches_2d(image, self.patch_size, self.max_patches, self.random_state)
}
pub fn summary(&self) -> String {
format!(
"PatchExtractor(patch_size={:?}, max_patches={:?}, strategy={:?})",
self.patch_size, self.max_patches, self.sampling_strategy
)
}
}
impl Default for PatchExtractor {
fn default() -> Self {
Self::new()
}
}
fn generate_patch_positions(
max_row: usize,
max_col: usize,
n_patches: usize,
max_patches: Option<usize>,
_random_state: Option<u64>,
) -> SklResult<Vec<(usize, usize)>> {
let total_patches = max_row * max_col;
let mut positions = Vec::new();
if let Some(max_p) = max_patches {
if max_p < total_patches {
let step = total_patches / max_p;
for i in (0..total_patches).step_by(step.max(1)) {
if positions.len() >= n_patches {
break;
}
let row = i / max_col;
let col = i % max_col;
positions.push((row, col));
}
} else {
for row in 0..max_row {
for col in 0..max_col {
if positions.len() >= n_patches {
break;
}
positions.push((row, col));
}
if positions.len() >= n_patches {
break;
}
}
}
} else {
for row in 0..max_row {
for col in 0..max_col {
positions.push((row, col));
}
}
}
Ok(positions)
}
fn extract_patches_optimized(
image: &ArrayView2<Float>,
patch_size: (usize, usize),
positions: &[(usize, usize)],
) -> SklResult<Array3<Float>> {
let (patch_height, patch_width) = patch_size;
let mut patches = Array3::zeros((positions.len(), patch_height, patch_width));
for (patch_idx, &(row, col)) in positions.iter().enumerate() {
for i in 0..patch_height {
for j in 0..patch_width {
patches[[patch_idx, i, j]] = image[[row + i, col + j]];
}
}
}
Ok(patches)
}
#[allow(non_snake_case)]
#[cfg(test)]
mod tests {
use super::*;
use scirs2_core::ndarray::Array2;
#[test]
fn test_extract_patches_2d() {
let image = Array2::from_shape_vec((4, 4), (0..16).map(|x| x as f64).collect())
.expect("operation should succeed");
let patches = extract_patches_2d(&image.view(), (2, 2), None, None)
.expect("operation should succeed");
assert_eq!(patches.dim(), (9, 2, 2)); assert_eq!(patches[[0, 0, 0]], 0.0); }
#[test]
fn test_extract_patches_max_limit() {
let image = Array2::from_shape_vec((6, 6), (0..36).map(|x| x as f64).collect())
.expect("operation should succeed");
let patches = extract_patches_2d(&image.view(), (2, 2), Some(5), None)
.expect("operation should succeed");
assert_eq!(patches.dim().0, 5); }
#[test]
fn test_patch_extractor_builder() {
let extractor = PatchExtractor::new()
.patch_size((3, 3))
.max_patches(Some(10))
.random_state(Some(42));
assert_eq!(extractor.patch_size, (3, 3));
assert_eq!(extractor.max_patches, Some(10));
assert_eq!(extractor.random_state, Some(42));
}
#[test]
fn test_reconstruct_from_patches() {
let original = Array2::from_shape_vec((4, 4), (0..16).map(|x| x as f64).collect())
.expect("operation should succeed");
let patches = extract_patches_2d(&original.view(), (2, 2), None, None)
.expect("operation should succeed");
let reconstructed =
reconstruct_from_patches_2d(&patches.view(), (4, 4)).expect("operation should succeed");
assert_eq!(reconstructed.dim(), (4, 4));
}
#[test]
fn test_invalid_patch_size() {
let image = Array2::from_shape_vec((2, 2), vec![1.0, 2.0, 3.0, 4.0])
.expect("operation should succeed");
let result = extract_patches_2d(&image.view(), (3, 3), None, None);
assert!(result.is_err());
}
#[test]
fn test_empty_image_reconstruction() {
let patches = Array3::zeros((0, 2, 2));
let result =
reconstruct_from_patches_2d(&patches.view(), (4, 4)).expect("operation should succeed");
assert_eq!(result.dim(), (4, 4));
assert!(result.iter().all(|&x| x == 0.0));
}
#[test]
fn test_patch_extractor_summary() {
let extractor = PatchExtractor::new().patch_size((5, 5));
let summary = extractor.summary();
assert!(summary.contains("patch_size=(5, 5)"));
}
#[test]
fn test_reconstruct_with_overlapping_patches() {
let image = Array2::from_shape_vec(
(4, 4),
vec![
1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0, 12.0, 13.0, 14.0, 15.0,
16.0,
],
)
.expect("operation should succeed");
let patches = extract_patches_2d(&image.view(), (2, 2), None, None)
.expect("operation should succeed");
assert_eq!(patches.dim(), (9, 2, 2));
let reconstructed =
reconstruct_from_patches_2d(&patches.view(), (4, 4)).expect("operation should succeed");
assert_eq!(reconstructed.dim(), (4, 4));
assert_eq!(reconstructed[[0, 0]], 1.0);
assert_eq!(reconstructed[[0, 3]], 4.0);
assert_eq!(reconstructed[[3, 0]], 13.0);
assert_eq!(reconstructed[[3, 3]], 16.0);
let expected_center = 6.0; assert!((reconstructed[[1, 1]] - expected_center).abs() < 1e-10);
}
#[test]
fn test_reconstruct_exact_coverage() {
let image = Array2::from_shape_vec((4, 4), (0..16).map(|x| x as f64).collect())
.expect("operation should succeed");
let patches = extract_patches_2d(&image.view(), (2, 2), None, None)
.expect("operation should succeed");
assert_eq!(patches.dim(), (9, 2, 2));
let reconstructed =
reconstruct_from_patches_2d(&patches.view(), (4, 4)).expect("operation should succeed");
assert_eq!(reconstructed.dim(), (4, 4));
assert_eq!(reconstructed[[0, 0]], image[[0, 0]]);
assert_eq!(reconstructed[[0, 3]], image[[0, 3]]);
assert_eq!(reconstructed[[3, 0]], image[[3, 0]]);
assert_eq!(reconstructed[[3, 3]], image[[3, 3]]);
}
#[test]
fn test_reconstruct_single_patch() {
let patch = Array3::from_shape_vec((1, 2, 2), vec![1.0, 2.0, 3.0, 4.0])
.expect("operation should succeed");
let reconstructed =
reconstruct_from_patches_2d(&patch.view(), (2, 2)).expect("operation should succeed");
assert_eq!(reconstructed.dim(), (2, 2));
assert_eq!(reconstructed[[0, 0]], 1.0);
assert_eq!(reconstructed[[0, 1]], 2.0);
assert_eq!(reconstructed[[1, 0]], 3.0);
assert_eq!(reconstructed[[1, 1]], 4.0);
}
}