use std::path::PathBuf;
use df_ocr_switcher::{DocPipeline, OcrModelPaths, OutputFormat, PpOcrEngine, PpOcrEngineConfig, TableModelPaths};
use ppocr_rs::PpOcrVersion;
fn main() {
if let Err(e) = run() {
eprintln!("[ocr-doc] {e}");
std::process::exit(1);
}
}
fn run() -> Result<(), Box<dyn std::error::Error>> {
let args: Vec<String> = std::env::args().collect();
if args.len() < 2 || args[1] == "--help" || args[1] == "-h" {
print_help();
return Ok(());
}
let input = PathBuf::from(&args[1]);
let mut format = OutputFormat::Markdown;
let mut engine_name = "ppocr".to_string();
let mut output_path: Option<PathBuf> = None;
let mut tier_str = "tiny".to_string();
let mut threads = 4usize;
let mut enable_tables = false;
let mut table_dir: Option<PathBuf> = None;
#[cfg(feature = "tesseract-engine")]
let mut tess_lang = "ita+eng".to_string();
let mut i = 2;
while i < args.len() {
match args[i].as_str() {
"--format" => {
i += 1;
format = match args.get(i).map(|s| s.as_str()) {
Some("json") => OutputFormat::Json,
_ => OutputFormat::Markdown,
};
}
"--engine" => {
i += 1;
engine_name = args.get(i).cloned().unwrap_or_default();
}
"--output" => {
i += 1;
output_path = args.get(i).map(PathBuf::from);
}
"--tier" => {
i += 1;
tier_str = args.get(i).cloned().unwrap_or("tiny".into());
}
"--threads" => {
i += 1;
threads = args.get(i).and_then(|s| s.parse().ok()).unwrap_or(4);
}
"--tables" => {
enable_tables = true;
}
"--table-dir" => {
i += 1;
table_dir = args.get(i).map(PathBuf::from);
}
#[cfg(feature = "tesseract-engine")]
"--lang" => {
i += 1;
tess_lang = args.get(i).cloned().unwrap_or("ita+eng".into());
}
_ => {}
}
i += 1;
}
if !input.exists() {
return Err(format!("file non trovato: {}", input.display()).into());
}
let layout_path = layout_model_path();
let tier = match tier_str.as_str() {
"small" => PpOcrVersion::V6Small,
"medium" => PpOcrVersion::V6Medium,
_ => PpOcrVersion::V6Tiny,
};
let pipeline_result: Result<DocPipeline, Box<dyn std::error::Error>> = match engine_name.as_str() {
#[cfg(feature = "tesseract-engine")]
"tesseract" => {
use df_ocr_switcher::{TesseractEngine, TesseractEngineConfig};
let ori_model = std::env::var("PPOCR_ORI_MODEL").ok().map(PathBuf::from);
let tessdata = std::env::var("TESSDATA_PREFIX").ok().map(PathBuf::from);
let engine = TesseractEngine::new(TesseractEngineConfig {
lang: tess_lang,
tessdata,
ori_model,
psm: None,
})?;
Ok(DocPipeline::new(Box::new(engine), &layout_path)?)
}
_ => {
let ori_model = std::env::var("PPOCR_ORI_MODEL").ok().map(PathBuf::from);
let ocr_models = ocr_models_from_env();
let tbl_models = if enable_tables { table_models_from_env(table_dir.as_deref()) } else { None };
let engine = PpOcrEngine::new(PpOcrEngineConfig {
tier,
ori_model,
ocr_models,
num_threads: threads,
table_models: tbl_models,
enable_tables,
enable_formula_decoder: false,
})?;
Ok(DocPipeline::new(Box::new(engine), &layout_path)?)
}
};
let mut pipeline = pipeline_result?;
let result = pipeline.process_file(&input, format)?;
match output_path {
Some(p) => std::fs::write(&p, &result)?,
None => print!("{result}"),
}
Ok(())
}
fn layout_model_path() -> PathBuf {
if let Ok(p) = std::env::var("PPOCR_LAYOUT_MODEL") {
return PathBuf::from(p);
}
let base = std::env::var("PPOCR_MODELS_DIR")
.unwrap_or_else(|_| "models/paddleocr".to_string());
PathBuf::from(base).join("layout").join("PP-DocLayoutV3.onnx")
}
fn ocr_models_from_env() -> Option<OcrModelPaths> {
let base = std::env::var("PPOCR_MODELS_DIR").ok()?;
let latin = PathBuf::from(base).join("latin");
let det = latin.join("det.onnx");
let rec = latin.join("rec_latin.onnx");
let dict = latin.join("dict_latin.txt");
if det.exists() && rec.exists() && dict.exists() {
Some(OcrModelPaths { det, rec, dict })
} else {
None
}
}
fn table_models_from_env(table_dir: Option<&std::path::Path>) -> Option<TableModelPaths> {
let dir = if let Some(d) = table_dir {
d.to_path_buf()
} else {
let base = std::env::var("PPOCR_MODELS_DIR").ok()?;
PathBuf::from(base).join("table")
};
let onnx = ["SLANet_plus.onnx", "SLANeXt_wired.onnx", "SLANet.onnx"]
.iter()
.map(|n| dir.join(n))
.find(|p| p.exists())?;
let dict = dir.join("table_structure_dict.txt");
if !dict.exists() {
eprintln!("[ocr-doc] table_structure_dict.txt non trovato in {}", dir.display());
return None;
}
let input_size = if onnx.file_name().and_then(|n| n.to_str())
.map(|n| n.starts_with("SLANet_plus")).unwrap_or(false) { Some(488) } else { None };
eprintln!("[ocr-doc] tabelle: {} (input {}px)",
onnx.file_name().unwrap().to_str().unwrap(),
input_size.unwrap_or(512));
Some(TableModelPaths { structure_onnx: onnx, structure_dict: dict, input_size })
}
fn print_help() {
eprintln!("ocr-doc <input> [OPTIONS]");
eprintln!();
eprintln!("Opzioni:");
eprintln!(" --format md|json output format (default: md)");
eprintln!(" --engine ppocr|tesseract engine OCR (default: ppocr)");
eprintln!(" --output <path> file output (default: stdout)");
eprintln!(" --tier tiny|small|medium PP-OCRv6 tier (default: tiny)");
eprintln!(" --threads <n> thread ONNX inference (default: 4)");
eprintln!(" --tables abilita table recognition (SLANet_plus)");
eprintln!(" --table-dir <path> dir modelli tabella (default: PPOCR_MODELS_DIR/table)");
eprintln!(" --lang <langs> lingue Tesseract (default: ita+eng)");
eprintln!();
eprintln!("Env vars:");
eprintln!(" PPOCR_LAYOUT_MODEL path PP-DocLayoutV3.onnx");
eprintln!(" PPOCR_MODELS_DIR dir base modelli (layout/, latin/, table/)");
eprintln!(" PPOCR_ORI_MODEL path orientation inference.onnx");
eprintln!(" ORT_DYLIB_PATH path onnxruntime.dll (ARM64)");
eprintln!(" TESSDATA_PREFIX tessdata dir (engine tesseract)");
eprintln!();
eprintln!("Esempi:");
eprintln!(" ocr-doc scan.tiff --tables");
eprintln!(" ocr-doc scan.tiff --tables --table-dir /models/table --format md");
eprintln!(" ocr-doc scan.tiff --engine tesseract --lang ita+eng");
}