ddddocr-rs/ddddocr-core/src/ocr/metadata.rs

//! OCR 模型元数据：归一化、缩放策略与模型信息。

// =====================================================================
// 1. 辅助定义的枚举与结构体
// =====================================================================

use crate::ocr::Charset;
use std::borrow::Cow;

/// 像素归一化策略。
#[derive(Debug, Clone, Copy)]
pub enum Normalization {
    /// 映射到 [0.0, 1.0] -> pixel / 255.0
    ZeroToOne,
    /// 映射到 [-1.0, 1.0] -> (pixel / 255.0 - 0.5) / 0.5
    MinusOneToOne,
}

impl Normalization {
    /// 对像素值执行归一化。
    #[inline(always)]
    pub fn normalize(&self, pixel: f32) -> f32 {
        match self {
            Normalization::ZeroToOne => pixel / 255.0,
            Normalization::MinusOneToOne => (pixel / 255.0 - 0.5) / 0.5,
        }
    }
}

/// 图像缩放策略枚举
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Resize {
    /// 固定宽高，例如 (64, 64)
    Fixed(u32, u32),
    /// 高度固定，宽度根据原始比例动态计算（对应 Python 的 [-1, H]）
    DynamicWidth(u32),
    /// 单字识别的正方形切图（对应 Python 的 word 为 True 且 [-1, H]）
    Square(u32),
}

/// OCR 模型元数据：字符集、缩放策略、通道数与归一化配置。
#[derive(Debug, Clone)]
pub struct ModelMetadata {
    /// 字符集。
    pub charset: Charset,
    /// 是否为单字识别模型
    pub word: bool,
    /// 缩放策略。
    pub resize: Resize,
    /// 图像通道数（1 或 3）
    pub channel: u8,
    /// 像素归一化配置。
    pub normalization: Normalization,
}

impl ModelMetadata {
    /// 创建模型元数据。
    pub fn new(
        charset: Charset,
        word: bool,
        resize: Resize,
        channel: u8,
        normalization: Normalization,
    ) -> Self {
        Self {
            charset,
            word,
            resize,
            channel,
            normalization,
        }
    }
    // --- 优雅的工厂模式构造器 ---
    /// 从静态字符切片构建元数据并自动生成字符集。
    pub fn from_static_slice(
        slice: &[&'static str],
        word: bool,
        resize: Resize,
        channel: u8,
        normalization: Normalization,
    ) -> Self {
        let tokens: Vec<Cow<'static, str>> = slice.iter().map(|&s| Cow::Borrowed(s)).collect();
        Self {
            charset: Charset::new(tokens),
            word,
            resize,
            channel,
            normalization,
        }
    }
    /// 使用内置旧版字符集（与官方旧版模型配套）。
    pub fn from_builtin_old(
        word: bool,
        resize: Resize,
        channel: u8,
        normalization: Normalization,
    ) -> Self {
        Self::from_static_slice(
            super::builtin::CHARSET_OLD,
            word,
            resize,
            channel,
            normalization,
        )
    }

    /// 使用内置 Beta 字符集（与官方 Beta 模型配套）。
    pub fn from_builtin_beta(
        word: bool,
        resize: Resize,
        channel: u8,
        normalization: Normalization,
    ) -> Self {
        Self::from_static_slice(
            super::builtin::CHARSET_BETA,
            word,
            resize,
            channel,
            normalization,
        )
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn normalization_zero_to_one() {
        let n = Normalization::ZeroToOne;
        assert_eq!(n.normalize(0.0), 0.0);
        assert_eq!(n.normalize(255.0), 1.0);
    }

    #[test]
    fn normalization_minus_one_to_one() {
        let n = Normalization::MinusOneToOne;
        assert_eq!(n.normalize(0.0), -1.0);
        assert_eq!(n.normalize(255.0), 1.0);
    }

    #[test]
    fn from_static_slice_builds_charset() {
        let meta = ModelMetadata::from_static_slice(
            &["", "a"],
            false,
            Resize::Fixed(64, 64),
            1,
            Normalization::ZeroToOne,
        );
        assert_eq!(meta.charset.size(), 2);
        assert_eq!(meta.charset.char_to_index("a"), 1);
        assert_eq!(meta.channel, 1);
        assert!(!meta.word);
    }
}