imageproc::find_contours panics on some edge images (called Option::unwrap() on None). Wrap with catch_unwind and provide a safe fallback (use full image bounds). Also add catch_unwind around the entire pipeline so no panic can crash the worker thread. Co-Authored-By: Kilo <kilo@kilo.ai>
198 lines
5.8 KiB
Rust
198 lines
5.8 KiB
Rust
use std::panic::catch_unwind;
|
|
|
|
use image::GrayImage;
|
|
use imageproc::contours::find_contours;
|
|
|
|
use tools_common::error::PipelineError;
|
|
|
|
/// Represents a detected corner point.
|
|
pub type CornerPoint = (f64, f64);
|
|
|
|
/// Find the 4 corners of the document from an edge image.
|
|
/// Wrapped in catch_unwind because imageproc's find_contours can panic.
|
|
pub fn detect_corners(edges: &GrayImage) -> Result<[CornerPoint; 4], FallbackReason> {
|
|
let result = catch_unwind(std::panic::AssertUnwindSafe(|| {
|
|
find_contours::<u8>(edges)
|
|
}));
|
|
|
|
let contours = match result {
|
|
Ok(c) => c,
|
|
Err(_) => return Err(FallbackReason::FindContoursPanic),
|
|
};
|
|
|
|
if contours.is_empty() {
|
|
return Err(FallbackReason::NoContours);
|
|
}
|
|
|
|
// Convert contours to use i32 coordinates
|
|
let contour_points: Vec<Vec<(i32, i32)>> = contours
|
|
.iter()
|
|
.map(|c| c.points.iter().map(|p| (p.x as i32, p.y as i32)).collect())
|
|
.collect();
|
|
|
|
// Sort by area descending
|
|
let mut sorted: Vec<_> = contour_points.iter().collect();
|
|
sorted.sort_by(|a, b| {
|
|
contour_area_slice(b)
|
|
.partial_cmp(&contour_area_slice(a))
|
|
.unwrap_or(std::cmp::Ordering::Equal)
|
|
});
|
|
|
|
for points in sorted.iter().take(5) {
|
|
let approx = catch_unwind(std::panic::AssertUnwindSafe(|| {
|
|
approx_quadrilateral(points)
|
|
}));
|
|
if let Ok(Some(corners)) = approx {
|
|
let ordered = order_corners(&corners);
|
|
return Ok(ordered);
|
|
}
|
|
}
|
|
|
|
// Fallback: use bounding rect of largest contour
|
|
if let Some(largest) = sorted.first() {
|
|
let rect = bounding_rect_slice(largest);
|
|
let corners = vec![
|
|
(rect.0 as f64, rect.1 as f64),
|
|
(rect.2 as f64, rect.1 as f64),
|
|
(rect.2 as f64, rect.3 as f64),
|
|
(rect.0 as f64, rect.3 as f64),
|
|
];
|
|
return Ok(order_corners(&corners));
|
|
}
|
|
|
|
Err(FallbackReason::NoContours)
|
|
}
|
|
|
|
/// The fallback reason if corner detection fails.
|
|
#[derive(Debug)]
|
|
pub enum FallbackReason {
|
|
FindContoursPanic,
|
|
NoContours,
|
|
NoRectangularContour,
|
|
TooSmall,
|
|
}
|
|
|
|
/// Compute the area of a contour using the Shoelace formula.
|
|
fn contour_area_slice(points: &[(i32, i32)]) -> f64 {
|
|
let n = points.len();
|
|
if n < 3 {
|
|
return 0.0;
|
|
}
|
|
let mut area = 0.0;
|
|
for i in 0..n {
|
|
let j = (i + 1) % n;
|
|
area += points[i].0 as f64 * points[j].1 as f64;
|
|
area -= points[j].0 as f64 * points[i].1 as f64;
|
|
}
|
|
area.abs() / 2.0
|
|
}
|
|
|
|
/// Approximate a contour to a quadrilateral using extreme points.
|
|
fn approx_quadrilateral(points: &[(i32, i32)]) -> Option<Vec<CornerPoint>> {
|
|
let n = points.len();
|
|
if n < 4 {
|
|
return None;
|
|
}
|
|
|
|
let top = points.iter().min_by(|a, b| a.1.cmp(&b.1))?;
|
|
let bottom = points.iter().max_by(|a, b| a.1.cmp(&b.1))?;
|
|
let left = points.iter().min_by(|a, b| a.0.cmp(&b.0))?;
|
|
let right = points.iter().max_by(|a, b| a.0.cmp(&b.0))?;
|
|
|
|
Some(vec![
|
|
(left.0 as f64, left.1 as f64),
|
|
(right.0 as f64, top.1 as f64),
|
|
(right.0 as f64, bottom.1 as f64),
|
|
(left.0 as f64, bottom.1 as f64),
|
|
])
|
|
}
|
|
|
|
/// Order 4 corners: top-left, top-right, bottom-right, bottom-left.
|
|
fn order_corners(points: &[CornerPoint]) -> [CornerPoint; 4] {
|
|
let mut pts: Vec<CornerPoint> = points.to_vec();
|
|
let mut ordered = [(0.0, 0.0); 4];
|
|
|
|
if pts.len() >= 4 {
|
|
pts.sort_by(|a, b| {
|
|
(a.0 + a.1)
|
|
.partial_cmp(&(b.0 + b.1))
|
|
.unwrap_or(std::cmp::Ordering::Equal)
|
|
});
|
|
ordered[0] = pts[0];
|
|
ordered[2] = pts[3];
|
|
|
|
pts.sort_by(|a, b| {
|
|
(a.0 - a.1)
|
|
.partial_cmp(&(b.0 - b.1))
|
|
.unwrap_or(std::cmp::Ordering::Equal)
|
|
});
|
|
ordered[1] = pts[3];
|
|
ordered[3] = pts[0];
|
|
}
|
|
|
|
ordered
|
|
}
|
|
|
|
/// Compute bounding rectangle: (left, top, right, bottom).
|
|
fn bounding_rect_slice(points: &[(i32, i32)]) -> (i32, i32, i32, i32) {
|
|
let left = points.iter().map(|p| p.0).min().unwrap_or(0);
|
|
let top = points.iter().map(|p| p.1).min().unwrap_or(0);
|
|
let right = points.iter().map(|p| p.0).max().unwrap_or(0);
|
|
let bottom = points.iter().map(|p| p.1).max().unwrap_or(0);
|
|
(left, top, right, bottom)
|
|
}
|
|
|
|
/// Detect corners with panic-safe fallback.
|
|
pub fn detect_corners_with_fallback(
|
|
edges: &GrayImage,
|
|
) -> Result<[CornerPoint; 4], PipelineError> {
|
|
if let Ok(corners) = detect_corners(edges) {
|
|
return Ok(corners);
|
|
}
|
|
|
|
// Attempt 2: half resolution
|
|
let (w, h) = (edges.width() / 2, edges.height() / 2);
|
|
if w > 10 && h > 10 {
|
|
let half = image::imageops::resize(
|
|
edges,
|
|
w,
|
|
h,
|
|
image::imageops::FilterType::Lanczos3,
|
|
);
|
|
if let Ok(corners) = detect_corners(&half) {
|
|
return Ok(corners.map(|(x, y)| (x * 2.0, y * 2.0)));
|
|
}
|
|
}
|
|
|
|
// Final fallback: use image bounds as corners (full image)
|
|
let (w, h) = (edges.width() as f64, edges.height() as f64);
|
|
tracing::warn!("Corner detection failed, using full image bounds");
|
|
Ok([(0.0, 0.0), (w, 0.0), (w, h), (0.0, h)])
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn test_contour_area_slice() {
|
|
let points = vec![(0, 0), (100, 0), (100, 100), (0, 100)];
|
|
let area = contour_area_slice(&points);
|
|
assert!((area - 10000.0).abs() < 1.0);
|
|
}
|
|
|
|
#[test]
|
|
fn test_bounding_rect_slice() {
|
|
let points = vec![(10, 20), (100, 30), (90, 150), (5, 140)];
|
|
let rect = bounding_rect_slice(&points);
|
|
assert_eq!(rect, (5, 20, 100, 150));
|
|
}
|
|
|
|
#[test]
|
|
fn test_order_corners() {
|
|
let pts = vec![(0.0, 100.0), (100.0, 100.0), (100.0, 0.0), (0.0, 0.0)];
|
|
let ordered = order_corners(&pts);
|
|
assert_eq!(ordered[0], (0.0, 0.0)); // TL
|
|
assert_eq!(ordered[2], (100.0, 100.0)); // BR
|
|
}
|
|
} |