use crate::core::error::Result;
use crate::core::gradient::GradientField;
use crate::core::image_view::OwnedImage;
use crate::core::scalar::Scalar;
use super::transform::AffineMap;
#[derive(Debug, Clone)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct AffineFrstConfig {
pub radius: u32,
pub gradient_threshold: Scalar,
pub smoothing_factor: Scalar,
pub affine_maps: Vec<AffineMap>,
}
impl Default for AffineFrstConfig {
fn default() -> Self {
Self {
radius: 8,
gradient_threshold: 0.0,
smoothing_factor: 0.5,
affine_maps: super::transform::sample_affine_maps(6, 3),
}
}
}
#[derive(Debug, Clone)]
pub struct AffineResponse {
pub response: OwnedImage<Scalar>,
pub affine_map: AffineMap,
pub peak_value: Scalar,
}
pub fn affine_frst_response_single(
gradient: &GradientField,
radius: u32,
affine: &AffineMap,
gradient_threshold: Scalar,
smoothing_factor: Scalar,
) -> Result<OwnedImage<Scalar>> {
let w = gradient.width();
let h = gradient.height();
let n = radius as Scalar;
let mut acc = OwnedImage::<Scalar>::zeros(w, h)?;
let acc_data = acc.data_mut();
let gx_data = gradient.gx.data();
let gy_data = gradient.gy.data();
let thresh_sq = gradient_threshold * gradient_threshold;
for y in 0..h {
for x in 0..w {
let idx = y * w + x;
let gx = gx_data[idx];
let gy = gy_data[idx];
let mag_sq = gx * gx + gy * gy;
if mag_sq < thresh_sq {
continue;
}
let mag = mag_sq.sqrt();
let dx = gx / mag;
let dy = gy / mag;
let wx = affine.a * dx + affine.b * dy;
let wy = affine.c * dx + affine.d * dy;
let wmag = (wx * wx + wy * wy).sqrt();
if wmag < 1e-8 {
continue;
}
let wx = wx / wmag;
let wy = wy / wmag;
let px = x as i32 + (wx * n).round() as i32;
let py = y as i32 + (wy * n).round() as i32;
if px >= 0 && (px as usize) < w && py >= 0 && (py as usize) < h {
acc_data[py as usize * w + px as usize] += mag;
}
}
}
let sigma = smoothing_factor * radius as Scalar;
if sigma > 0.5 {
crate::core::blur::gaussian_blur_inplace(&mut acc, sigma);
}
Ok(acc)
}
pub fn affine_frst_responses(
gradient: &GradientField,
config: &AffineFrstConfig,
) -> Result<Vec<AffineResponse>> {
let mut responses = Vec::with_capacity(config.affine_maps.len());
for affine in &config.affine_maps {
let response = affine_frst_response_single(
gradient,
config.radius,
affine,
config.gradient_threshold,
config.smoothing_factor,
)?;
let peak_value = response.data().iter().copied().fold(0.0f32, Scalar::max);
responses.push(AffineResponse {
response,
affine_map: *affine,
peak_value,
});
}
responses.sort_by(|a, b| b.peak_value.partial_cmp(&a.peak_value).unwrap());
Ok(responses)
}
#[cfg(test)]
mod tests {
use super::*;
use crate::core::gradient::sobel_gradient;
use crate::core::image_view::ImageView;
use crate::core::nms::{non_maximum_suppression, NmsConfig};
fn make_ellipse_image(size: usize, cx: f32, cy: f32, a: f32, b: f32, angle: f32) -> Vec<u8> {
let cos_a = angle.cos();
let sin_a = angle.sin();
let mut data = vec![0u8; size * size];
for y in 0..size {
for x in 0..size {
let dx = x as f32 - cx;
let dy = y as f32 - cy;
let lx = dx * cos_a + dy * sin_a;
let ly = -dx * sin_a + dy * cos_a;
if (lx / a).powi(2) + (ly / b).powi(2) <= 1.0 {
data[y * size + x] = 255;
}
}
}
data
}
#[test]
fn affine_responses_sorted_by_peak() {
let size = 80;
let data = make_ellipse_image(size, 40.0, 40.0, 12.0, 8.0, 0.3);
let image = ImageView::from_slice(&data, size, size).unwrap();
let grad = sobel_gradient(&image).unwrap();
let config = AffineFrstConfig {
radius: 10,
affine_maps: super::super::transform::sample_affine_maps(4, 2),
..AffineFrstConfig::default()
};
let responses = affine_frst_responses(&grad, &config).unwrap();
assert_eq!(responses.len(), 8);
for pair in responses.windows(2) {
assert!(pair[0].peak_value >= pair[1].peak_value);
}
}
#[test]
fn best_affine_finds_ellipse_center() {
let size = 80;
let cx = 40.0;
let cy = 40.0;
let data = make_ellipse_image(size, cx, cy, 12.0, 8.0, 0.0);
let image = ImageView::from_slice(&data, size, size).unwrap();
let grad = sobel_gradient(&image).unwrap();
let config = AffineFrstConfig {
radius: 10,
..AffineFrstConfig::default()
};
let responses = affine_frst_responses(&grad, &config).unwrap();
assert!(!responses.is_empty());
let best = &responses[0];
let peaks = non_maximum_suppression(
&best.response.view(),
&NmsConfig {
radius: 5,
threshold: 0.0,
max_detections: 1,
},
);
assert!(!peaks.is_empty(), "should find a peak");
let err = ((peaks[0].position.x - cx).powi(2) + (peaks[0].position.y - cy).powi(2)).sqrt();
assert!(
err < 8.0,
"peak at ({}, {}) too far from center ({cx}, {cy}), error={err}",
peaks[0].position.x,
peaks[0].position.y
);
}
}