use crate::core::buffer::DocBuffer;
use crate::edge::morphology::Morphology;
use crate::edge::threshold::AdaptiveThreshold;
use crate::error::DocQuadError;
use fast_canny::{CannyConfig, CannyWorkspace, canny_u8};
use std::time::Instant;
pub struct EdgeDetector {
workspace: CannyWorkspace,
width: usize,
height: usize,
}
impl EdgeDetector {
pub fn new(width: usize, height: usize) -> Result<Self, DocQuadError> {
let workspace = CannyWorkspace::new(width, height).map_err(|e| {
log::error!(
"[Edge::Detector] - Failed to create CannyWorkspace: {:?}",
e
);
DocQuadError::EdgeDetectionError
})?;
log::info!(
"[Edge::Detector] - CannyWorkspace pre-allocated for {}x{}.",
width,
height
);
Ok(Self {
workspace,
width,
height,
})
}
pub fn detect(&mut self, buffer: &DocBuffer<'_>) -> Result<Vec<u8>, DocQuadError> {
self.detect_with_debug(buffer, None)
}
pub fn detect_with_debug(
&mut self,
buffer: &DocBuffer<'_>,
debug_dir: Option<&std::path::Path>,
) -> Result<Vec<u8>, DocQuadError> {
let start = Instant::now();
let view = buffer.as_array_view()?;
let (low, high) = AdaptiveThreshold::calculate(&view, 0.33);
let mut hist = [0u32; 256];
for &pixel in view.iter() {
hist[pixel as usize] += 1;
}
let total_pixels = (buffer.width * buffer.height) as f32;
let dark_pct = hist[0..64].iter().sum::<u32>() as f32 / total_pixels * 100.0;
let mid_pct = hist[64..192].iter().sum::<u32>() as f32 / total_pixels * 100.0;
let bright_pct = hist[192..256].iter().sum::<u32>() as f32 / total_pixels * 100.0;
log::info!(
"[Edge::Detector] - Pixel brightness distribution: \
dark(0-63)={:.1}%, mid(64-191)={:.1}%, bright(192-255)={:.1}%",
dark_pct, mid_pct, bright_pct
);
let mut max_bucket_count = 0u32;
let mut max_bucket_start = 0usize;
for i in (0..256).step_by(16) {
let bucket_sum: u32 = hist[i..i + 16].iter().sum();
if bucket_sum > max_bucket_count {
max_bucket_count = bucket_sum;
max_bucket_start = i;
}
}
log::info!(
"[Edge::Detector] - Histogram peak bucket: [{}-{}] = {} pixels ({:.1}%)",
max_bucket_start,
max_bucket_start + 15,
max_bucket_count,
max_bucket_count as f32 / total_pixels * 100.0
);
log::info!(
"[Edge::Detector] - Adaptive thresholds: low={:.2}, high={:.2} \
(image={}x{}, sigma=0.33)",
low, high, buffer.width, buffer.height
);
if high > 200.0 {
log::warn!(
"[Edge::Detector] - Canny high threshold {:.2} is very high (>200). \
Document edges with weak gradients may be missed.",
high
);
}
let input_data: Vec<u8> = if buffer.stride == buffer.width {
log::debug!(
"[Edge::Detector] - Input is contiguous (stride==width={}), direct copy.",
buffer.width
);
buffer.data[..(buffer.width * buffer.height) as usize].to_vec()
} else {
log::debug!(
"[Edge::Detector] - Input has stride padding (stride={} > width={}), compacting rows.",
buffer.stride, buffer.width
);
let mut compact = Vec::with_capacity((buffer.width * buffer.height) as usize);
for row in view.rows() {
compact.extend(row.iter().copied());
}
compact
};
if let Some(dir) = debug_dir {
Self::save_gray_image(
&input_data,
self.width,
self.height,
&dir.join("debug_01_input_gray.png"),
);
}
#[cfg(debug_assertions)]
{
log::info!(
"[Edge::Detector] - [DEBUG] Threshold sensitivity scan \
(sigma=1.0, 3 trial configs):"
);
for &(trial_low, trial_high) in &[(2.0f32, 8.0), (5.0, 20.0), (10.0, 40.0)] {
if let Ok(cfg) = CannyConfig::builder()
.sigma(1.0)
.thresholds(trial_low, trial_high)
.build()
{
if let Ok(trial_slice) = canny_u8(&input_data, &mut self.workspace, &cfg) {
let trial_count = trial_slice.iter().filter(|&&v| v == 255).count();
let trial_density =
trial_count as f32 / (self.width * self.height) as f32 * 100.0;
log::info!(
"[Edge::Detector] - [DEBUG] low={:.1}, high={:.1} -> \
edge_pixels={}, density={:.2}%",
trial_low, trial_high, trial_count, trial_density
);
}
}
}
log::info!(
"[Edge::Detector] - [DEBUG] Low-sigma scan (sigma=0.5, low=5.0, high=20.0):"
);
if let Ok(cfg_low_sigma) = CannyConfig::builder()
.sigma(0.5)
.thresholds(5.0, 20.0)
.build()
{
if let Ok(trial_slice) = canny_u8(&input_data, &mut self.workspace, &cfg_low_sigma) {
let trial_count = trial_slice.iter().filter(|&&v| v == 255).count();
let trial_density =
trial_count as f32 / (self.width * self.height) as f32 * 100.0;
log::info!(
"[Edge::Detector] - [DEBUG] sigma=0.5, low=5.0, high=20.0 -> \
edge_pixels={}, density={:.2}%",
trial_count, trial_density
);
}
}
}
let cfg = CannyConfig::builder()
.sigma(1.0)
.thresholds(low, high)
.build()
.map_err(|e| {
log::error!("[Edge::Detector] - Invalid Canny config: {:?}", e);
DocQuadError::EdgeDetectionError
})?;
let canny_start = Instant::now();
let edge_slice = canny_u8(&input_data, &mut self.workspace, &cfg).map_err(|e| {
log::error!("[Edge::Detector] - canny_u8 failed: {:?}", e);
DocQuadError::EdgeDetectionError
})?;
let raw_edges = edge_slice.to_vec();
let canny_elapsed = canny_start.elapsed().as_millis();
let raw_edge_count = raw_edges.iter().filter(|&&v| v == 255).count();
let raw_density = raw_edge_count as f32 / (self.width * self.height) as f32 * 100.0;
log::info!(
"[Edge::Detector] - Canny raw result: edge_pixels={}, density={:.2}%, \
size={}x{}. Canny elapsed: {}ms",
raw_edge_count, raw_density, self.width, self.height, canny_elapsed
);
#[cfg(debug_assertions)]
Self::log_edge_spatial_distribution(&raw_edges, self.width, self.height, "raw_canny");
if let Some(dir) = debug_dir {
Self::save_binary_image(
&raw_edges,
self.width,
self.height,
&dir.join("debug_02_canny_raw.png"),
);
}
let morph_radius = Self::choose_morph_radius(self.width, self.height);
log::debug!(
"[Edge::Detector] - Applying morphological close: radius={} (image={}x{})",
morph_radius, self.width, self.height
);
let closed_edges = Morphology::close(&raw_edges, self.width, self.height, morph_radius);
let closed_edge_count = closed_edges.iter().filter(|&&v| v == 255).count();
let closed_density =
closed_edge_count as f32 / (self.width * self.height) as f32 * 100.0;
let net_change = closed_edge_count as i64 - raw_edge_count as i64;
let growth_pct = if raw_edge_count > 0 {
net_change as f32 / raw_edge_count as f32 * 100.0
} else {
0.0
};
log::info!(
"[Edge::Detector] - After morphological close: edge_pixels={}, density={:.2}%, \
net_change={:+} ({:+.1}%). Total elapsed: {}ms",
closed_edge_count, closed_density, net_change, growth_pct,
start.elapsed().as_millis()
);
#[cfg(debug_assertions)]
Self::log_edge_spatial_distribution(
&closed_edges,
self.width,
self.height,
"after_close",
);
if let Some(dir) = debug_dir {
Self::save_binary_image(
&closed_edges,
self.width,
self.height,
&dir.join("debug_03_after_close.png"),
);
#[cfg(debug_assertions)]
{
let closed_r2 = Morphology::close(&raw_edges, self.width, self.height, 2);
let r2_count = closed_r2.iter().filter(|&&v| v == 255).count();
log::info!(
"[Edge::Detector] - [DEBUG] radius=2 close result: \
edge_pixels={}, density={:.2}%",
r2_count,
r2_count as f32 / (self.width * self.height) as f32 * 100.0
);
Self::save_binary_image(
&closed_r2,
self.width,
self.height,
&dir.join("debug_04_close_radius2.png"),
);
let closed_r3 = Morphology::close(&raw_edges, self.width, self.height, 3);
let r3_count = closed_r3.iter().filter(|&&v| v == 255).count();
log::info!(
"[Edge::Detector] - [DEBUG] radius=3 close result: \
edge_pixels={}, density={:.2}%",
r3_count,
r3_count as f32 / (self.width * self.height) as f32 * 100.0
);
Self::save_binary_image(
&closed_r3,
self.width,
self.height,
&dir.join("debug_05_close_radius3.png"),
);
if let Ok(cfg_s05) = CannyConfig::builder()
.sigma(0.5)
.thresholds(5.0, 20.0)
.build()
{
if let Ok(s05_slice) = canny_u8(&input_data, &mut self.workspace, &cfg_s05) {
let s05_raw = s05_slice.to_vec();
let s05_closed = Morphology::close(&s05_raw, self.width, self.height, 2);
let s05_count = s05_closed.iter().filter(|&&v| v == 255).count();
log::info!(
"[Edge::Detector] - [DEBUG] sigma=0.5 + radius=2 close: \
edge_pixels={}, density={:.2}%",
s05_count,
s05_count as f32 / (self.width * self.height) as f32 * 100.0
);
Self::save_binary_image(
&s05_closed,
self.width,
self.height,
&dir.join("debug_06_sigma05_close_r2.png"),
);
}
}
}
}
if closed_density > 25.0 {
log::warn!(
"[Edge::Detector] - Post-close edge density {:.2}% is very high (>25%). \
Morphological close may have over-connected noise. \
Consider reducing morph_radius ({}) or raising Canny thresholds.",
closed_density, morph_radius
);
} else if closed_density < 0.05 {
log::warn!(
"[Edge::Detector] - Post-close edge density {:.2}% is very low (<0.05%). \
Document edges may be missing. \
Consider lowering Canny thresholds (current low={:.2}, high={:.2}).",
closed_density, low, high
);
}
if growth_pct > 100.0 {
log::warn!(
"[Edge::Detector] - Morphological close growth rate {:.1}% is very high (>100%). \
radius={} may be too large, causing noise over-connection.",
growth_pct, morph_radius
);
}
Ok(closed_edges)
}
fn log_edge_spatial_distribution(
edges: &[u8],
width: usize,
height: usize,
label: &str,
) {
let grid_cols = 4usize;
let grid_rows = 4usize;
let cell_w = width / grid_cols;
let cell_h = height / grid_rows;
log::info!(
"[Edge::Detector] - [DEBUG] Edge spatial distribution ({}) \
in {}x{} grid (cell={}x{}px):",
label, grid_cols, grid_rows, cell_w, cell_h
);
let cell_total = (cell_w * cell_h) as f32;
let mut grid_lines = Vec::new();
for gy in 0..grid_rows {
let mut row_str = String::from(" |");
for gx in 0..grid_cols {
let x0 = gx * cell_w;
let y0 = gy * cell_h;
let x1 = ((gx + 1) * cell_w).min(width);
let y1 = ((gy + 1) * cell_h).min(height);
let mut count = 0usize;
for y in y0..y1 {
for x in x0..x1 {
if edges[y * width + x] == 255 {
count += 1;
}
}
}
let density = count as f32 / cell_total * 100.0;
let symbol = if density < 1.0 {
" "
} else if density < 5.0 {
"."
} else if density < 15.0 {
"o"
} else if density < 30.0 {
"O"
} else {
"#"
};
row_str.push_str(&format!("{:>5.1}%{}|", density, symbol));
}
grid_lines.push(row_str);
}
let separator = format!(" +{}+", "------+".repeat(grid_cols));
log::info!("[Edge::Detector] - [DEBUG] {}", separator);
for line in &grid_lines {
log::info!("[Edge::Detector] - [DEBUG] {}", line);
}
log::info!("[Edge::Detector] - [DEBUG] {}", separator);
}
fn save_gray_image(data: &[u8], width: usize, height: usize, path: &std::path::Path) {
match Self::write_pgm(data, width, height, path.with_extension("pgm").as_path()) {
Ok(_) => log::info!(
"[Edge::Detector] - [DEBUG] Saved gray image: {}",
path.with_extension("pgm").display()
),
Err(e) => log::warn!(
"[Edge::Detector] - [DEBUG] Failed to save gray image {}: {}",
path.display(),
e
),
}
}
fn save_binary_image(data: &[u8], width: usize, height: usize, path: &std::path::Path) {
match Self::write_pgm(data, width, height, path.with_extension("pgm").as_path()) {
Ok(_) => log::info!(
"[Edge::Detector] - [DEBUG] Saved binary image: {}",
path.with_extension("pgm").display()
),
Err(e) => log::warn!(
"[Edge::Detector] - [DEBUG] Failed to save binary image {}: {}",
path.display(),
e
),
}
}
fn write_pgm(
data: &[u8],
width: usize,
height: usize,
path: &std::path::Path,
) -> std::io::Result<()> {
use std::io::Write;
if let Some(parent) = path.parent() {
std::fs::create_dir_all(parent)?;
}
let mut file = std::fs::File::create(path)?;
write!(file, "P5\n{} {}\n255\n", width, height)?;
file.write_all(&data[..width * height])?;
Ok(())
}
fn choose_morph_radius(width: usize, height: usize) -> usize {
let long_edge = width.max(height);
if long_edge <= 512 { 1 }
else if long_edge <= 1024 { 2 } else { 3 }
}
}