a3s_use_ocr/
assets.rs

1use std::path::{Path, PathBuf};
2
3use a3s_use_core::{UseError, UseResult};
4use serde::{Deserialize, Serialize};
5
6use crate::config::{load_detection, load_recognition, MODEL_FAMILY};
7
8pub(crate) const RECEIPT_FILE: &str = ".a3s-ppocr-v6.json";
9
10#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
11#[serde(rename_all = "kebab-case")]
12pub enum OcrInstallSource {
13    Environment,
14    Packaged,
15    Managed,
16    Missing,
17}
18
19#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
20#[serde(rename_all = "camelCase")]
21pub struct OcrRuntimeStatus {
22    pub available: bool,
23    pub source: OcrInstallSource,
24    pub model: String,
25    pub model_dir: Option<PathBuf>,
26    pub managed_root: Option<PathBuf>,
27    pub detail: String,
28}
29
30#[derive(Debug, Clone)]
31pub(crate) struct ModelAssets {
32    pub(crate) root: PathBuf,
33    pub(crate) detection_model: PathBuf,
34    pub(crate) detection_config: PathBuf,
35    pub(crate) recognition_model: PathBuf,
36    pub(crate) recognition_config: PathBuf,
37    pub(crate) source: OcrInstallSource,
38}
39
40pub fn ocr_status() -> OcrRuntimeStatus {
41    let managed_root = managed_root().ok();
42    match resolve_model_assets() {
43        Ok(assets) => OcrRuntimeStatus {
44            available: true,
45            source: assets.source,
46            model: MODEL_FAMILY.to_string(),
47            model_dir: Some(assets.root),
48            managed_root,
49            detail: "ready".to_string(),
50        },
51        Err(error) => OcrRuntimeStatus {
52            available: false,
53            source: error
54                .details
55                .get("source")
56                .and_then(serde_json::Value::as_str)
57                .map(source_from_name)
58                .unwrap_or(OcrInstallSource::Missing),
59            model: MODEL_FAMILY.to_string(),
60            model_dir: error
61                .details
62                .get("modelDir")
63                .and_then(serde_json::Value::as_str)
64                .map(PathBuf::from),
65            managed_root,
66            detail: error.message,
67        },
68    }
69}
70
71pub(crate) fn resolve_model_assets() -> UseResult<ModelAssets> {
72    if let Some(path) = std::env::var_os("A3S_OCR_MODEL_DIR")
73        .filter(|value| !value.is_empty())
74        .map(PathBuf::from)
75    {
76        let path = absolute(path)?;
77        return validate_assets(&path, OcrInstallSource::Environment);
78    }
79
80    let managed = managed_model_dir()?;
81    if path_exists(&managed)? {
82        return validate_assets(&managed, OcrInstallSource::Managed);
83    }
84
85    if let Ok(executable) = std::env::current_exe() {
86        if let Some(parent) = executable.parent() {
87            let packaged = parent.join("ocr-models").join(MODEL_FAMILY);
88            if path_exists(&packaged)? {
89                return validate_assets(&packaged, OcrInstallSource::Packaged);
90            }
91        }
92    }
93
94    Err(UseError::new(
95        "use.ocr.model_missing",
96        format!("The local {MODEL_FAMILY} model bundle is not installed."),
97    )
98    .with_suggestion("Run 'a3s install use/ocr'.")
99    .with_detail("source", "missing")
100    .with_detail("modelDir", managed.display().to_string()))
101}
102
103pub(crate) fn validate_assets(root: &Path, source: OcrInstallSource) -> UseResult<ModelAssets> {
104    let root = std::fs::canonicalize(root).map_err(|error| {
105        model_error(
106            source,
107            root,
108            format!(
109                "Failed to resolve the {MODEL_FAMILY} model directory '{}': {error}",
110                root.display()
111            ),
112        )
113    })?;
114    let detection_model = checked_file(&root, "det/inference.onnx", 256 * 1024 * 1024, source)?;
115    let detection_config = checked_file(&root, "det/inference.yml", 2 * 1024 * 1024, source)?;
116    let recognition_model = checked_file(&root, "rec/inference.onnx", 256 * 1024 * 1024, source)?;
117    let recognition_config = checked_file(&root, "rec/inference.yml", 2 * 1024 * 1024, source)?;
118
119    load_detection(&detection_config)?;
120    load_recognition(&recognition_config)?;
121
122    Ok(ModelAssets {
123        root,
124        detection_model,
125        detection_config,
126        recognition_model,
127        recognition_config,
128        source,
129    })
130}
131
132pub(crate) fn managed_root() -> UseResult<PathBuf> {
133    if let Some(value) = std::env::var_os("A3S_USE_OCR_HOME") {
134        return absolute(PathBuf::from(value));
135    }
136    if let Some(value) = std::env::var_os("A3S_DATA_HOME") {
137        return Ok(absolute(PathBuf::from(value))?.join("use/ocr"));
138    }
139    if let Some(value) = std::env::var_os("XDG_DATA_HOME") {
140        return Ok(absolute(PathBuf::from(value))?.join("a3s/use/ocr"));
141    }
142    if let Some(home) = std::env::var_os("HOME").map(PathBuf::from) {
143        return Ok(absolute(home)?.join(".local/share/a3s/use/ocr"));
144    }
145    #[cfg(windows)]
146    if let Some(value) = std::env::var_os("LOCALAPPDATA") {
147        return Ok(absolute(PathBuf::from(value))?.join("a3s/use/ocr"));
148    }
149    Err(UseError::new(
150        "use.ocr.data_home_missing",
151        "Cannot determine the A3S Use OCR data directory.",
152    ))
153}
154
155pub(crate) fn managed_model_dir() -> UseResult<PathBuf> {
156    Ok(managed_root()?.join(MODEL_FAMILY))
157}
158
159fn checked_file(
160    root: &Path,
161    relative: &str,
162    max_bytes: u64,
163    source: OcrInstallSource,
164) -> UseResult<PathBuf> {
165    let path = root.join(relative);
166    let canonical = std::fs::canonicalize(&path).map_err(|error| {
167        model_error(
168            source,
169            root,
170            format!(
171                "Required {MODEL_FAMILY} asset '{}' is unreadable: {error}",
172                path.display()
173            ),
174        )
175    })?;
176    if !canonical.starts_with(root) {
177        return Err(model_error(
178            source,
179            root,
180            format!(
181                "Required {MODEL_FAMILY} asset '{}' escapes its model directory.",
182                path.display()
183            ),
184        ));
185    }
186    let metadata = std::fs::metadata(&canonical).map_err(|error| {
187        model_error(
188            source,
189            root,
190            format!(
191                "Failed to inspect {MODEL_FAMILY} asset '{}': {error}",
192                canonical.display()
193            ),
194        )
195    })?;
196    if !metadata.is_file() || metadata.len() == 0 || metadata.len() > max_bytes {
197        return Err(model_error(
198            source,
199            root,
200            format!(
201                "{MODEL_FAMILY} asset '{}' must be a non-empty regular file no larger than {max_bytes} bytes.",
202                canonical.display()
203            ),
204        ));
205    }
206    Ok(canonical)
207}
208
209fn path_exists(path: &Path) -> UseResult<bool> {
210    match std::fs::symlink_metadata(path) {
211        Ok(_) => Ok(true),
212        Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(false),
213        Err(error) => Err(UseError::new(
214            "use.ocr.model_unreadable",
215            format!(
216                "Failed to inspect OCR model path '{}': {error}",
217                path.display()
218            ),
219        )),
220    }
221}
222
223fn model_error(source: OcrInstallSource, root: &Path, message: impl Into<String>) -> UseError {
224    UseError::new("use.ocr.model_invalid", message)
225        .with_suggestion("Run 'a3s install use/ocr --force' to restore the pinned PP-OCRv6 bundle.")
226        .with_detail("source", source_name(source))
227        .with_detail("modelDir", root.display().to_string())
228}
229
230fn source_name(source: OcrInstallSource) -> &'static str {
231    match source {
232        OcrInstallSource::Environment => "environment",
233        OcrInstallSource::Packaged => "packaged",
234        OcrInstallSource::Managed => "managed",
235        OcrInstallSource::Missing => "missing",
236    }
237}
238
239fn source_from_name(value: &str) -> OcrInstallSource {
240    match value {
241        "environment" => OcrInstallSource::Environment,
242        "packaged" => OcrInstallSource::Packaged,
243        "managed" => OcrInstallSource::Managed,
244        _ => OcrInstallSource::Missing,
245    }
246}
247
248fn absolute(path: PathBuf) -> UseResult<PathBuf> {
249    if path.is_absolute() {
250        Ok(path)
251    } else {
252        std::env::current_dir()
253            .map(|directory| directory.join(path))
254            .map_err(|error| {
255                UseError::new(
256                    "use.ocr.path_resolution_failed",
257                    format!("Failed to resolve OCR data path: {error}"),
258                )
259            })
260    }
261}
a3s_use_ocr/assets.rs

a3s_use_ocr/
assets.rs
