visual_cortex_ocr_onnx/
models.rs

1use std::path::{Path, PathBuf};
2use std::time::Duration;
3
4use sha2::{Digest, Sha256};
5
6use crate::error::OcrError;
7
8/// Generous ceiling for a ~10 MB model download; surfaces hung connections
9/// as Download errors instead of blocking forever.
10const DOWNLOAD_TIMEOUT: Duration = Duration::from_secs(300);
11
12/// One pinned, checksummed model artifact.
13#[derive(Debug, Clone, Copy)]
14pub struct ModelFile {
15    pub filename: &'static str,
16    pub url: &'static str,
17    pub sha256: &'static str,
18}
19
20/// PP-OCRv4 mobile detection model (multilingual; finds text boxes).
21pub const DET_MODEL: ModelFile = ModelFile {
22    filename: "ch_PP-OCRv4_det_infer.onnx",
23    url: "https://huggingface.co/SWHL/RapidOCR/resolve/main/PP-OCRv4/ch_PP-OCRv4_det_infer.onnx",
24    sha256: "d2a7720d45a54257208b1e13e36a8479894cb74155a5efe29462512d42f49da9",
25};
26
27/// PP-OCRv3 English mobile recognition model (reads each box).
28pub const REC_MODEL: ModelFile = ModelFile {
29    filename: "en_PP-OCRv3_rec_infer.onnx",
30    url: "https://huggingface.co/SWHL/RapidOCR/resolve/main/PP-OCRv3/en_PP-OCRv3_rec_infer.onnx",
31    sha256: "ef7abd8bd3629ae57ea2c28b425c1bd258a871b93fd2fe7c433946ade9b5d9ea",
32};
33
34/// English recognition charset (95 entries; CTC classes = blank + 95 + space).
35pub const REC_DICT: ModelFile = ModelFile {
36    filename: "en_dict.txt",
37    url: "https://raw.githubusercontent.com/PaddlePaddle/PaddleOCR/v2.7.0/ppocr/utils/en_dict.txt",
38    sha256: "5662df9d2d03f0e8ca0d3b0649d6acbab904b6a14b3d3521463c71c37c668ce3",
39};
40
41pub(crate) fn sha256_hex(bytes: &[u8]) -> String {
42    let digest = Sha256::digest(bytes);
43    digest.iter().map(|b| format!("{b:02x}")).collect()
44}
45
46/// The default cache directory: `<platform cache dir>/visual-cortex`.
47pub(crate) fn default_cache_dir() -> PathBuf {
48    dirs::cache_dir()
49        .unwrap_or_else(std::env::temp_dir)
50        .join("visual-cortex")
51}
52
53/// Return the cached path for `file`, downloading and verifying if needed.
54/// A cached file failing its checksum is deleted and re-downloaded once.
55pub(crate) async fn ensure_cached(file: &ModelFile, dir: &Path) -> Result<PathBuf, OcrError> {
56    let path = dir.join(file.filename);
57    if path.exists() {
58        let bytes = tokio::fs::read(&path)
59            .await
60            .map_err(|e| OcrError::Download {
61                url: file.url.to_string(),
62                reason: format!("cache read failed: {e}"),
63            })?;
64        if sha256_hex(&bytes) == file.sha256 {
65            return Ok(path);
66        }
67        tracing::warn!(
68            "checksum mismatch for cached {}; re-downloading",
69            path.display()
70        );
71        let _ = tokio::fs::remove_file(&path).await;
72    }
73
74    tokio::fs::create_dir_all(dir)
75        .await
76        .map_err(|e| OcrError::Download {
77            url: file.url.to_string(),
78            reason: format!("cache dir creation failed: {e}"),
79        })?;
80    let client = reqwest::Client::builder()
81        .timeout(DOWNLOAD_TIMEOUT)
82        .build()
83        .map_err(|e| OcrError::Download {
84            url: file.url.to_string(),
85            reason: format!("http client build failed: {e}"),
86        })?;
87    let response = client
88        .get(file.url)
89        .send()
90        .await
91        .map_err(|e| OcrError::Download {
92            url: file.url.to_string(),
93            reason: e.to_string(),
94        })?;
95    if !response.status().is_success() {
96        return Err(OcrError::Download {
97            url: file.url.to_string(),
98            reason: format!("HTTP {}", response.status()),
99        });
100    }
101    let bytes = response.bytes().await.map_err(|e| OcrError::Download {
102        url: file.url.to_string(),
103        reason: e.to_string(),
104    })?;
105    let actual = sha256_hex(&bytes);
106    if actual != file.sha256 {
107        return Err(OcrError::ChecksumMismatch {
108            path: path.display().to_string(),
109            expected: file.sha256.to_string(),
110            actual,
111        });
112    }
113    // Write via a temp sibling + atomic rename so a concurrent reader can
114    // never observe a partially-written model file.
115    let tmp = path.with_extension("tmp");
116    tokio::fs::write(&tmp, &bytes)
117        .await
118        .map_err(|e| OcrError::Download {
119            url: file.url.to_string(),
120            reason: format!("cache write failed: {e}"),
121        })?;
122    tokio::fs::rename(&tmp, &path)
123        .await
124        .map_err(|e| OcrError::Download {
125            url: file.url.to_string(),
126            reason: format!("cache rename failed: {e}"),
127        })?;
128    Ok(path)
129}
130
131/// Ensure all three artifacts are cached; returns (det, rec, dict) paths.
132pub async fn ensure_all_cached() -> Result<(PathBuf, PathBuf, PathBuf), OcrError> {
133    let dir = default_cache_dir();
134    Ok((
135        ensure_cached(&DET_MODEL, &dir).await?,
136        ensure_cached(&REC_MODEL, &dir).await?,
137        ensure_cached(&REC_DICT, &dir).await?,
138    ))
139}
140
141#[cfg(test)]
142mod tests {
143    use super::*;
144
145    const HELLO_SHA256: &str = "2cf24dba5fb0a30e26e83b2ac5b9e29e1b161e5c1fa7425e73043362938b9824";
146
147    #[test]
148    fn sha256_hex_computes_known_digest() {
149        assert_eq!(sha256_hex(b"hello"), HELLO_SHA256);
150    }
151
152    #[tokio::test]
153    async fn ensure_cached_accepts_preseeded_valid_file() {
154        let dir = tempfile::tempdir().unwrap();
155        let file = ModelFile {
156            filename: "hello.bin",
157            // Bogus URL: the network must never be touched when the cache is valid.
158            url: "https://invalid.invalid/hello.bin",
159            sha256: HELLO_SHA256,
160        };
161        std::fs::write(dir.path().join("hello.bin"), b"hello").unwrap();
162        let path = ensure_cached(&file, dir.path()).await.unwrap();
163        assert_eq!(path, dir.path().join("hello.bin"));
164    }
165
166    #[tokio::test]
167    async fn ensure_cached_deletes_corrupt_file_and_errors() {
168        let dir = tempfile::tempdir().unwrap();
169        let file = ModelFile {
170            filename: "hello.bin",
171            // Port 1 on loopback refuses the connection immediately -- no DNS
172            // dependency, so the test stays hermetic and fast everywhere.
173            url: "http://127.0.0.1:1/hello.bin",
174            sha256: HELLO_SHA256,
175        };
176        std::fs::write(dir.path().join("hello.bin"), b"corrupted").unwrap();
177        // Corrupt cache -> deleted -> re-download attempted -> Download error
178        // (the connection is refused by construction).
179        let err = ensure_cached(&file, dir.path()).await.unwrap_err();
180        assert!(matches!(err, OcrError::Download { .. }), "got {err:?}");
181        assert!(
182            !dir.path().join("hello.bin").exists(),
183            "corrupt file must be deleted"
184        );
185    }
186}
visual_cortex_ocr_onnx/models.rs

visual_cortex_ocr_onnx/
models.rs
