rapid-structure-rs 1.0.0

use std::{path::PathBuf, time::Instant};

use ndarray::{ArrayView4, Axis, s};
use serde::{Deserialize, Serialize};

use crate::{
    Quad,
    config::{LangDet, ModelType, OcrVersion, RecImage, RuntimeConfig},
    error::{RapidOcrError, Result},
    model_registry::ModelRegistry,
    model_store::{
        DownloadRequest, default_model_store_dir, download_verified, require_model_hash,
        verify_existing_file,
    },
    model_verify::VerifiedModel,
    ocr::session::{OcrSession, OcrSessionKind},
    runtime::provider::ProviderResolution,
};

use super::{
    postprocess::DbPostProcess,
    preprocess::{DetPreProcess, DetPreprocessScratch},
};

#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(default, deny_unknown_fields)]
pub struct DetectorConfig {
    pub lang: LangDet,
    pub ocr_version: OcrVersion,
    pub model_type: ModelType,
    pub model_path: Option<PathBuf>,
    pub allow_download: bool,
    pub limit_side_len: usize,
    pub limit_type: String,
    pub std: [f32; 3],
    pub mean: [f32; 3],
    pub thresh: f32,
    pub box_thresh: f32,
    pub max_candidates: usize,
    pub unclip_ratio: f32,
    pub use_dilation: bool,
    pub score_mode: String,
    pub model_store_dir: Option<PathBuf>,
}

impl Default for DetectorConfig {
    fn default() -> Self {
        Self {
            lang: LangDet::Ch,
            ocr_version: OcrVersion::PPocrV4,
            model_type: ModelType::Mobile,
            model_path: None,
            allow_download: true,
            limit_side_len: 736,
            limit_type: "min".to_string(),
            std: [0.5, 0.5, 0.5],
            mean: [0.5, 0.5, 0.5],
            thresh: 0.3,
            box_thresh: 0.5,
            max_candidates: 1000,
            unclip_ratio: 1.6,
            use_dilation: true,
            score_mode: "fast".to_string(),
            model_store_dir: None,
        }
    }
}

#[derive(Debug, Clone, Default)]
pub struct DetOutput {
    pub boxes: Vec<Quad>,
    pub scores: Vec<f32>,
    pub elapsed_ms: f32,
    pub breakdown: Option<DetTimingBreakdown>,
    /// Actual spatial shape (width, height) submitted to the detector ORT session.
    pub input_size: (u32, u32),
}

#[derive(Debug, Clone, Copy, Default)]
pub struct DetTimingBreakdown {
    pub preprocess_ms: f32,
    pub infer_ms: f32,
    pub postprocess_ms: f32,
}

#[derive(Debug)]
pub struct Detector {
    pre: DetPreProcess,
    post: DbPostProcess,
    session: OcrSession,
    batch_scratch: Vec<f32>,
    preprocess_scratch: DetPreprocessScratch,
}

impl Detector {
    /// `runtime` 来自引擎唯一的运行时档案（`RuntimeProfile::session_runtime`）：
    /// 检测阶段的配置里刻意没有线程/provider 字段，避免出现第二个设置来源。
    pub fn new(config: DetectorConfig, runtime: &RuntimeConfig) -> Result<Self> {
        let model_store_dir = config
            .model_store_dir
            .clone()
            .unwrap_or_else(default_model_store_dir);

        let model_path = if let Some(path) = &config.model_path {
            verify_existing_file(path)?
        } else if config.allow_download {
            let registry = ModelRegistry::from_default_yaml()?;
            let resolved =
                registry.resolve_det(config.ocr_version, config.lang, config.model_type)?;
            // §6.4：哈希必填。默认表的 det 条目都记录了 SHA-256，缺项会在这里变成
            // 可定位错误，而不是"下载一个不校验的文件"。
            let expected = require_model_hash(resolved.sha256.as_deref(), &resolved.model_url)?;
            download_verified(&DownloadRequest::new(
                &resolved.model_url,
                expected,
                &model_store_dir,
            ))?
        } else {
            return Err(RapidOcrError::Config(
                "detector model_path is not set and allow_download=false".to_string(),
            ));
        };

        let pre = DetPreProcess {
            limit_side_len: config.limit_side_len,
            limit_type: config.limit_type,
            mean: config.mean,
            std: config.std,
        };
        let post = DbPostProcess {
            thresh: config.thresh,
            box_thresh: config.box_thresh,
            max_candidates: config.max_candidates,
            unclip_ratio: config.unclip_ratio,
            use_dilation: config.use_dilation,
            score_mode: config.score_mode,
            ..DbPostProcess::default()
        };
        // 文本模型的完整性在**准入/下载层**已保证（`verify_existing_file` /
        // `download_verified` 都强制哈希）；会话层拿不到部署声明的期望摘要，因此
        // 构造“无期望”凭证（仍检查存在性），结构上不再存在“忘记校验”的旁路。
        let model = VerifiedModel::verify(&model_path, None)?;
        let session = OcrSession::new(&model, runtime, OcrSessionKind::Det)?;
        Ok(Self {
            pre,
            post,
            session,
            batch_scratch: Vec::new(),
            preprocess_scratch: DetPreprocessScratch::default(),
        })
    }

    pub fn detect(&mut self, img: &RecImage) -> Result<DetOutput> {
        let start = Instant::now();
        // `DBPostProcess` expects destination size in the detector input image space
        // (before detector-side resize), matching RapidOCR Python behavior.
        let det_limit_side_len = resolve_limit_side_len_like_python(
            self.pre.limit_type.as_str(),
            self.pre.limit_side_len,
            img.width().max(img.height()),
        );
        let pre_start = Instant::now();
        let (resized_h, resized_w) = self.pre.run_into_buffer_with_scratch(
            img,
            &mut self.batch_scratch,
            &mut self.preprocess_scratch,
            Some(det_limit_side_len),
        )?;
        let preprocess_ms = pre_start.elapsed().as_secs_f32() * 1000.0;
        let batch_view = ArrayView4::from_shape(
            (1, 3, resized_h, resized_w),
            &self.batch_scratch[..3 * resized_h * resized_w],
        )
        .map_err(|e| RapidOcrError::InvalidInput(format!("invalid det batch shape: {e}")))?;
        let post = &self.post;
        crate::runtime::memory::tl_probe("detect_infer");
        let infer_start = Instant::now();
        let mut postprocess_ms = 0.0_f32;
        let (boxes, scores) = self.session.run_array4_view_with(batch_view, |preds| {
            if preds.len_of(Axis(0)) == 0 || preds.len_of(Axis(1)) == 0 {
                return Ok((Vec::new(), Vec::new()));
            }
            let post_start = Instant::now();
            let map = preds.slice(s![0, 0, .., ..]);
            let out = post.run_view(map, img.width(), img.height());
            postprocess_ms = post_start.elapsed().as_secs_f32() * 1000.0;
            Ok(out)
        })?;
        let infer_total_ms = infer_start.elapsed().as_secs_f32() * 1000.0;
        let infer_ms = (infer_total_ms - postprocess_ms).max(0.0);
        crate::runtime::memory::tl_probe("detect_post");

        Ok(DetOutput {
            boxes,
            scores,
            elapsed_ms: start.elapsed().as_secs_f32() * 1000.0,
            breakdown: Some(DetTimingBreakdown {
                preprocess_ms,
                infer_ms,
                postprocess_ms,
            }),
            input_size: (resized_w as u32, resized_h as u32),
        })
    }

    pub fn update_postprocess(&mut self, box_thresh: Option<f32>, unclip_ratio: Option<f32>) {
        if let Some(v) = box_thresh {
            self.post.box_thresh = v;
        }
        if let Some(v) = unclip_ratio {
            self.post.unclip_ratio = v;
        }
    }

    pub fn provider_resolution(&self) -> ProviderResolution {
        self.session.provider_resolution()
    }

    pub fn model_path(&self) -> &std::path::Path {
        self.session.model_path()
    }
}

pub(crate) fn resolve_limit_side_len_like_python(
    limit_type: &str,
    configured_limit_side_len: usize,
    max_wh: usize,
) -> usize {
    if limit_type == "min" {
        configured_limit_side_len
    } else if max_wh < 960 {
        960
    } else if max_wh < 1500 {
        1500
    } else {
        2000
    }
}

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

    #[test]
    fn limit_side_len_for_min_uses_configured_value() {
        assert_eq!(resolve_limit_side_len_like_python("min", 736, 2048), 736);
    }

    #[test]
    fn limit_side_len_for_non_min_matches_python_buckets() {
        assert_eq!(resolve_limit_side_len_like_python("max", 736, 800), 960);
        assert_eq!(resolve_limit_side_len_like_python("max", 736, 1200), 1500);
        assert_eq!(resolve_limit_side_len_like_python("max", 736, 1800), 2000);
    }
}