fast_canny/

lib.rs

pub mod gaussian;
pub mod hysteresis;
pub mod kernel;
pub mod nms;
pub mod pipeline;
pub mod workspace;

// src/lib.rs
#[cfg(target_arch = "x86_64")]
pub mod avx2_kernel {
    pub use crate::kernel::avx2::*;
}

mod pipeline_ptr;

pub use workspace::CannyWorkspace;

use gaussian::apply as gaussian_apply;
use hysteresis::track_edges;
use pipeline::execute_tiled_pipeline;

// =====================================================================
// CannyConfig
// =====================================================================

#[derive(Debug, Clone)]
pub struct CannyConfig {
    pub sigma: f32,
    pub low_thresh: f32,
    pub high_thresh: f32,
}

impl Default for CannyConfig {
    fn default() -> Self {
        Self {
            sigma: 1.0,
            low_thresh: 50.0,
            high_thresh: 150.0,
        }
    }
}

impl CannyConfig {
    pub fn builder() -> CannyConfigBuilder {
        CannyConfigBuilder::default()
    }
}

#[derive(Default)]
pub struct CannyConfigBuilder {
    sigma: Option<f32>,
    low_thresh: Option<f32>,
    high_thresh: Option<f32>,
}

impl CannyConfigBuilder {
    pub fn sigma(mut self, v: f32) -> Self {
        self.sigma = Some(v);
        self
    }
    pub fn thresholds(mut self, low: f32, high: f32) -> Self {
        self.low_thresh = Some(low);
        self.high_thresh = Some(high);
        self
    }
    pub fn build(self) -> Result<CannyConfig, CannyError> {
        let cfg = CannyConfig {
            sigma: self.sigma.unwrap_or(1.0),
            low_thresh: self.low_thresh.unwrap_or(50.0),
            high_thresh: self.high_thresh.unwrap_or(150.0),
        };
        if cfg.low_thresh > cfg.high_thresh {
            return Err(CannyError::InvalidThresholds {
                low: cfg.low_thresh,
                high: cfg.high_thresh,
            });
        }
        Ok(cfg)
    }
}

// =====================================================================
// CannyError / CannyStatus
// =====================================================================

#[derive(Debug, Clone, PartialEq)]
pub enum CannyError {
    InvalidDimensions { width: usize, height: usize },
    InputLengthMismatch { expected: usize, actual: usize },
    InvalidThresholds { low: f32, high: f32 },
    NullPointer,
}

impl std::fmt::Display for CannyError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::InvalidDimensions { width, height } => {
                write!(f, "invalid dimensions: {}x{} (min 3x3)", width, height)
            }
            Self::InputLengthMismatch { expected, actual } => {
                write!(f, "input length: expected {}, got {}", expected, actual)
            }
            Self::InvalidThresholds { low, high } => {
                write!(f, "thresholds: low={} > high={}", low, high)
            }
            Self::NullPointer => write!(f, "null pointer"),
        }
    }
}

impl std::error::Error for CannyError {}

#[repr(i32)]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CannyStatus {
    Ok = 0,
    NullPointer = -1,
    InvalidDimensions = -2,
    InputLengthMismatch = -3,
    InvalidThresholds = -4,
}

impl From<&CannyError> for CannyStatus {
    fn from(e: &CannyError) -> Self {
        match e {
            CannyError::NullPointer => Self::NullPointer,
            CannyError::InvalidDimensions { .. } => Self::InvalidDimensions,
            CannyError::InputLengthMismatch { .. } => Self::InputLengthMismatch,
            CannyError::InvalidThresholds { .. } => Self::InvalidThresholds,
        }
    }
}

// =====================================================================
// canny()：接受 f32 切片的主 API（零拷贝路径）
// =====================================================================

pub fn canny<'ws>(
    src: &[f32],
    ws: &'ws mut CannyWorkspace,
    cfg: &CannyConfig,
) -> Result<&'ws [u8], CannyError> {
    let expected = ws.width * ws.height;
    if src.len() != expected {
        return Err(CannyError::InputLengthMismatch {
            expected,
            actual: src.len(),
        });
    }
    if cfg.low_thresh > cfg.high_thresh {
        return Err(CannyError::InvalidThresholds {
            low: cfg.low_thresh,
            high: cfg.high_thresh,
        });
    }

    ws.reset();

    let w = ws.width;
    let h = ws.height;

    // 高斯平滑：写入 buffer_b，或直接使用 src（零拷贝）
    if cfg.sigma > 0.0 {
        gaussian_apply(src, &mut ws.buffer_b, w, h, cfg.sigma);
        // 取裸指针后重新构造切片，脱离对 ws.buffer_b 的借用，
        // 避免后续 execute_tiled_pipeline 同时可变借用 ws 时产生冲突
        let blurred_ptr = ws.buffer_b.as_ptr();
        let blurred_len = ws.buffer_b.len();
        let blurred: &[f32] = unsafe { std::slice::from_raw_parts(blurred_ptr, blurred_len) };
        execute_tiled_pipeline(blurred, ws, cfg.low_thresh, cfg.high_thresh);
    } else {
        execute_tiled_pipeline(src, ws, cfg.low_thresh, cfg.high_thresh);
    }

    track_edges(&mut ws.edge_map, w, h, &ws.arena);

    Ok(&ws.edge_map)
}

// =====================================================================
// canny_u8()：新增 API，直接接受 &[u8] 输入
//
// 设计说明：
//   u8 → f32 转换不可避免（Sobel 需要浮点运算），
//   但转换结果写入 ws.buffer_b，后续所有阶段均零拷贝操作 workspace 内缓冲区。
//   相比让调用方自行转换，此 API 封装了转换细节，提供更友好的接口。
// =====================================================================

/// 接受 u8 灰度图的 Canny 边缘检测 API。
///
/// # 参数
/// - `src`：输入灰度图，每像素 1 字节，长度必须等于 `ws.width * ws.height`
/// - `ws`：预分配的工作区，可跨帧复用
/// - `cfg`：算法配置（sigma、低/高阈值）
///
/// # 返回值
/// `&[u8]` 二值边缘图，仅含 0（非边缘）和 255（边缘），
/// 生命周期绑定到 `ws`，下次调用前有效。
///
/// # 零拷贝说明
/// u8→f32 转换写入 `ws.buffer_b`（预分配，无堆分配），
/// 后续 Sobel/NMS/Hysteresis 均直接操作 workspace 内缓冲区。
pub fn canny_u8<'ws>(
    src: &[u8],
    ws: &'ws mut CannyWorkspace,
    cfg: &CannyConfig,
) -> Result<&'ws [u8], CannyError> {
    let expected = ws.width * ws.height;
    if src.len() != expected {
        return Err(CannyError::InputLengthMismatch {
            expected,
            actual: src.len(),
        });
    }
    if cfg.low_thresh > cfg.high_thresh {
        return Err(CannyError::InvalidThresholds {
            low: cfg.low_thresh,
            high: cfg.high_thresh,
        });
    }

    ws.reset();

    let w = ws.width;
    let h = ws.height;

    // u8 → f32 转换，写入预分配的 buffer_b（无堆分配）
    // 值域保持 [0.0, 255.0]，与 f32 API 语义一致
    {
        let dst = &mut ws.buffer_b;
        for (d, &s) in dst.iter_mut().zip(src.iter()) {
            *d = s as f32;
        }
    }

    // 根据 sigma 决定是否再次高斯平滑
    if cfg.sigma > 0.0 {
        // 就地平滑：从 buffer_b 读，写回 buffer_b
        // 需要临时借用，使用 buffer_a 作为中间缓冲
        let src_f32: Vec<f32> = ws.buffer_b.clone();
        gaussian_apply(&src_f32, &mut ws.buffer_b, w, h, cfg.sigma);
    }

    // 此时 buffer_b 已包含最终的平滑输入，直接传引用（零拷贝）
    let blurred = ws.buffer_b.as_slice();

    // SAFETY: blurred 是 ws.buffer_b 的切片，execute_tiled_pipeline 内部
    // 通过裸指针访问 ws 的其他字段，不与 buffer_b 产生别名冲突
    // （buffer_b 只读，其他缓冲区只写）
    let blurred_ptr = blurred.as_ptr();
    let blurred_len = blurred.len();
    let blurred_static: &'static [f32] = unsafe {
        // SAFETY: 生命周期延长仅用于传递给 execute_tiled_pipeline，
        // 该函数在返回前不会 drop ws，blurred 在整个调用期间有效
        std::slice::from_raw_parts(blurred_ptr, blurred_len)
    };

    execute_tiled_pipeline(blurred_static, ws, cfg.low_thresh, cfg.high_thresh);
    track_edges(&mut ws.edge_map, w, h, &ws.arena);

    Ok(&ws.edge_map)
}

// =====================================================================
// C-FFI
// =====================================================================

#[unsafe(no_mangle)]
pub extern "C" fn canny_workspace_new(width: usize, height: usize) -> *mut CannyWorkspace {
    match CannyWorkspace::new(width, height) {
        Ok(ws) => Box::into_raw(Box::new(ws)),
        Err(e) => {
            log::error!("[canny_workspace_new] {}", e);
            std::ptr::null_mut()
        }
    }
}

#[unsafe(no_mangle)]
pub unsafe extern "C" fn canny_workspace_free(ws: *mut CannyWorkspace) {
    if !ws.is_null() {
        // SAFETY: ws 由 canny_workspace_new 通过 Box::into_raw 创建
        drop(unsafe { Box::from_raw(ws) });
    }
}

#[unsafe(no_mangle)]
pub unsafe extern "C" fn canny_process_ex(
    ws: *mut CannyWorkspace,
    src: *const f32,
    src_len: usize,
    sigma: f32,
    low_thresh: f32,
    high_thresh: f32,
    out_edge: *mut *const u8,
) -> i32 {
    if ws.is_null() || src.is_null() || out_edge.is_null() {
        log::error!("[canny_process_ex] null pointer");
        return CannyStatus::NullPointer as i32;
    }

    // SAFETY: ws 由调用方保证有效；src 由调用方保证长度为 src_len
    let ws_ref = unsafe { &mut *ws };
    let src_slice = unsafe { std::slice::from_raw_parts(src, src_len) };

    let cfg = match CannyConfig::builder()
        .sigma(sigma)
        .thresholds(low_thresh, high_thresh)
        .build()
    {
        Ok(c) => c,
        Err(e) => {
            log::error!("[canny_process_ex] config: {}", e);
            return CannyStatus::from(&e) as i32;
        }
    };

    match canny(src_slice, ws_ref, &cfg) {
        Ok(edge) => {
            // SAFETY: out_edge 非空（上方已检查）
            unsafe {
                *out_edge = edge.as_ptr();
            }
            CannyStatus::Ok as i32
        }
        Err(e) => {
            log::error!("[canny_process_ex] {}", e);
            CannyStatus::from(&e) as i32
        }
    }
}

/// C-FFI：接受 u8 输入的 Canny 接口
#[unsafe(no_mangle)]
pub unsafe extern "C" fn canny_process_u8(
    ws: *mut CannyWorkspace,
    src: *const u8,
    src_len: usize,
    sigma: f32,
    low_thresh: f32,
    high_thresh: f32,
    out_edge: *mut *const u8,
) -> i32 {
    if ws.is_null() || src.is_null() || out_edge.is_null() {
        log::error!("[canny_process_u8] null pointer");
        return CannyStatus::NullPointer as i32;
    }

    // SAFETY: 指针由调用方保证有效
    let ws_ref = unsafe { &mut *ws };
    let src_slice = unsafe { std::slice::from_raw_parts(src, src_len) };

    let cfg = match CannyConfig::builder()
        .sigma(sigma)
        .thresholds(low_thresh, high_thresh)
        .build()
    {
        Ok(c) => c,
        Err(e) => {
            log::error!("[canny_process_u8] config: {}", e);
            return CannyStatus::from(&e) as i32;
        }
    };

    match canny_u8(src_slice, ws_ref, &cfg) {
        Ok(edge) => {
            unsafe {
                *out_edge = edge.as_ptr();
            }
            CannyStatus::Ok as i32
        }
        Err(e) => {
            log::error!("[canny_process_u8] {}", e);
            CannyStatus::from(&e) as i32
        }
    }
}