fast-canny 0.1.0

//! 图像边缘检测示例
//!
//! 用法：
//!   cargo run --example detect_image -- <输入路径> [输出路径] [sigma] [low_thresh] [high_thresh]
//!
//! 示例：
//!   cargo run --example detect_image -- input.jpg
//!   cargo run --example detect_image -- input.png output_edge.png 1.0 50.0 150.0

use fast_canny::{CannyConfig, CannyWorkspace, canny};
use image::{DynamicImage, GrayImage, ImageError, Luma};
use std::path::{Path, PathBuf};
use std::time::Instant;

// =====================================================================
// 错误类型
// =====================================================================

#[derive(Debug)]
enum DetectError {
    /// 图片文件不存在
    FileNotFound(PathBuf),
    /// 不支持的图片格式
    UnsupportedFormat(String),
    /// 图片尺寸不满足最小要求（< 3×3）
    ImageTooSmall { width: u32, height: u32 },
    /// 图片 I/O 错误
    ImageIo(ImageError),
    /// Canny 算法错误
    Canny(fast_canny::CannyError),
}

impl std::fmt::Display for DetectError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::FileNotFound(p) => write!(f, "file not found: {}", p.display()),
            Self::UnsupportedFormat(ext) => {
                write!(
                    f,
                    "unsupported format: .{ext} (supported: png, jpg, jpeg, bmp, tiff, webp)"
                )
            }
            Self::ImageTooSmall { width, height } => {
                write!(f, "image {width}x{height} is too small, minimum is 3x3")
            }
            Self::ImageIo(e) => write!(f, "image I/O error: {e}"),
            Self::Canny(e) => write!(f, "canny error: {e}"),
        }
    }
}

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

impl From<ImageError> for DetectError {
    fn from(e: ImageError) -> Self {
        Self::ImageIo(e)
    }
}

impl From<fast_canny::CannyError> for DetectError {
    fn from(e: fast_canny::CannyError) -> Self {
        Self::Canny(e)
    }
}

// =====================================================================
// 支持的图片格式白名单
// =====================================================================

const SUPPORTED_EXTENSIONS: &[&str] = &["png", "jpg", "jpeg", "bmp", "tiff", "tif", "webp"];

fn check_format(path: &Path) -> Result<(), DetectError> {
    let ext = path
        .extension()
        .and_then(|e| e.to_str())
        .unwrap_or("")
        .to_lowercase();

    if SUPPORTED_EXTENSIONS.contains(&ext.as_str()) {
        Ok(())
    } else {
        Err(DetectError::UnsupportedFormat(ext))
    }
}

// =====================================================================
// 图片加载：转换为 f32 单通道灰度图
// =====================================================================

/// 加载图片并转换为 f32 灰度数组（值域 [0.0, 255.0]）
fn load_image_as_f32(path: &Path) -> Result<(Vec<f32>, u32, u32), DetectError> {
    if !path.exists() {
        return Err(DetectError::FileNotFound(path.to_path_buf()));
    }

    check_format(path)?;

    let t = Instant::now();
    let img = image::open(path)?;

    let gray: GrayImage = img.into_luma8();
    let (width, height) = gray.dimensions();

    log::info!(
        "[load_image] loaded {} — {}x{}, elapsed={:?}",
        path.display(),
        width,
        height,
        t.elapsed()
    );

    if width < 3 || height < 3 {
        return Err(DetectError::ImageTooSmall { width, height });
    }

    // u8 → f32，值域保持 [0.0, 255.0]
    let pixels: Vec<f32> = gray.pixels().map(|Luma([v])| *v as f32).collect();

    Ok((pixels, width, height))
}

// =====================================================================
// 边缘图保存：u8 二值图 → PNG/JPG
// =====================================================================

/// 将 Canny 输出的 u8 边缘图保存为图片文件
fn save_edge_map(
    edge_map: &[u8],
    width: u32,
    height: u32,
    output_path: &Path,
) -> Result<(), DetectError> {
    check_format(output_path)?;

    // 构造灰度图：0 → 黑，255 → 白
    let gray = GrayImage::from_raw(width, height, edge_map.to_vec())
        .expect("edge_map size does not match width x height");

    let t = Instant::now();
    DynamicImage::ImageLuma8(gray).save(output_path)?;

    log::info!(
        "[save_edge_map] saved {} — elapsed={:?}",
        output_path.display(),
        t.elapsed()
    );

    Ok(())
}

// =====================================================================
// 生成默认输出路径
//
// 规则：<原文件名>_edge.<原扩展名>
// 例如：photo.jpg → photo_edge.jpg
// =====================================================================

fn default_output_path(input: &Path) -> PathBuf {
    let stem = input
        .file_stem()
        .and_then(|s| s.to_str())
        .unwrap_or("output");
    let ext = input.extension().and_then(|e| e.to_str()).unwrap_or("png");
    let filename = format!("{stem}_edge.{ext}");

    input
        .parent()
        .unwrap_or_else(|| Path::new("."))
        .join(filename)
}

// =====================================================================
// 命令行参数解析
// =====================================================================

struct Args {
    input: PathBuf,
    output: PathBuf,
    sigma: f32,
    low_thresh: f32,
    high_thresh: f32,
}

fn parse_args() -> Result<Args, String> {
    let args: Vec<String> = std::env::args().collect();

    if args.len() < 2 {
        return Err(format!(
            "usage: {} <input> [output] [sigma] [low_thresh] [high_thresh]\n\
             example: {} input.jpg output_edge.png 1.0 50.0 150.0",
            args[0], args[0]
        ));
    }

    let input = PathBuf::from(&args[1]);
    // 未指定输出路径时，自动生成 <stem>_edge.<ext>
    let output = args
        .get(2)
        .map(PathBuf::from)
        .unwrap_or_else(|| default_output_path(&input));

    let sigma: f32 = args.get(3).and_then(|s| s.parse().ok()).unwrap_or(1.0);

    let low_thresh: f32 = args.get(4).and_then(|s| s.parse().ok()).unwrap_or(50.0);

    let high_thresh: f32 = args.get(5).and_then(|s| s.parse().ok()).unwrap_or(150.0);

    if low_thresh > high_thresh {
        return Err(format!(
            "invalid args: low_thresh ({low_thresh}) must be <= high_thresh ({high_thresh})"
        ));
    }

    Ok(Args {
        input,
        output,
        sigma,
        low_thresh,
        high_thresh,
    })
}

// =====================================================================
// 主流程
// =====================================================================

fn run() -> Result<(), DetectError> {
    // ── 初始化日志 ──────────────────────────────────────────────────
    env_logger::builder()
        .filter_level(log::LevelFilter::Info)
        .init();

    // ── 解析参数 ────────────────────────────────────────────────────
    let args = parse_args().unwrap_or_else(|e| {
        eprintln!("error: {e}");
        std::process::exit(1);
    });

    println!("=== fast-canny edge detection ===");
    println!("input:       {}", args.input.display());
    println!("output:      {}", args.output.display());
    println!("sigma:       {}", args.sigma);
    println!("low_thresh:  {}", args.low_thresh);
    println!("high_thresh: {}", args.high_thresh);
    println!();

    // ── 步骤 1：加载图片 ────────────────────────────────────────────
    let t_total = Instant::now();

    let (src_pixels, width, height) = load_image_as_f32(&args.input)?;
    println!(
        "✓ image loaded: {}x{} ({} pixels)",
        width,
        height,
        src_pixels.len()
    );

    // ── 步骤 2：创建 Workspace ──────────────────────────────────────
    let t_ws = Instant::now();
    let mut ws = CannyWorkspace::new(width as usize, height as usize)?;
    log::info!(
        "[workspace] created {}x{}, elapsed={:.1}ms",
        width,
        height,
        t_ws.elapsed().as_secs_f64() * 1000.0
    );
    println!(
        "✓ workspace created ({:.1} ms)",
        t_ws.elapsed().as_secs_f64() * 1000.0
    );

    // ── 步骤 3：配置 Canny 参数 ─────────────────────────────────────
    let cfg = CannyConfig::builder()
        .sigma(args.sigma)
        .thresholds(args.low_thresh, args.high_thresh)
        .build()
        .map_err(DetectError::Canny)?;

    // ── 步骤 4：执行 Canny 边缘检测 ────────────────────────────────
    let t_canny = Instant::now();
    let edge_map = canny(&src_pixels, &mut ws, &cfg)?;
    let canny_ms = t_canny.elapsed().as_secs_f64() * 1000.0;

    let edge_count = edge_map.iter().filter(|&&v| v == 255).count();
    let edge_ratio = edge_count as f64 / edge_map.len() as f64 * 100.0;

    log::info!(
        "[canny] done — edge_pixels={} ({:.2}%), elapsed={:.2}ms",
        edge_count,
        edge_ratio,
        canny_ms
    );
    println!(
        "✓ canny done ({:.2} ms) — edge pixels: {} ({:.2}%)",
        canny_ms, edge_count, edge_ratio
    );

    // ── 步骤 5：保存结果 ────────────────────────────────────────────
    // 确保输出目录存在
    if let Some(parent) = args.output.parent() {
        if !parent.as_os_str().is_empty() && !parent.exists() {
            std::fs::create_dir_all(parent)
                .map_err(|e| DetectError::ImageIo(ImageError::IoError(e)))?;
        }
    }

    save_edge_map(edge_map, width, height, &args.output)?;
    println!("✓ result saved: {}", args.output.display());

    println!();
    println!(
        "total elapsed: {:.2} ms",
        t_total.elapsed().as_secs_f64() * 1000.0
    );

    log::info!(
        "[detect_image] finished — input={}, output={}, total={:.2}ms",
        args.input.display(),
        args.output.display(),
        t_total.elapsed().as_secs_f64() * 1000.0
    );

    Ok(())
}

fn main() {
    if let Err(e) = run() {
        eprintln!("error: {e}");
        std::process::exit(1);
    }
}