rusto-rs 0.1.6

RustO! - Pure Rust OCR library based on RapidOCR with PaddleOCR engine
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146

use clap::{Parser, ValueEnum};
use rusto::rusto_ocr::RustO as RustOCR;
use rusto::{OrientClassifier, OrientConfig};
use std::path::PathBuf;

#[derive(Parser)]
#[command(name = "rusto")]
#[command(about = "RustO! - Pure Rust OCR powered by PaddleOCR engine", long_about = None)]
struct Cli {
    /// Path to detection model (ONNX)
    #[arg(long)]
    det_model: PathBuf,

    /// Path to recognition model (ONNX)
    #[arg(long)]
    rec_model: PathBuf,

    /// Path to dictionary file
    #[arg(long)]
    dict: PathBuf,

    /// Input image path
    image: PathBuf,

    /// Output format
    #[arg(short, long, value_enum, default_value_t = OutputFormat::Json)]
    format: OutputFormat,

    /// Enable document orientation classification
    #[arg(long)]
    enable_orient: bool,

    /// Path to orientation classification model (lcnet.mnn)
    #[arg(long)]
    orient_model: Option<PathBuf>,

    /// Enable text unwarping
    #[arg(long)]
    enable_unwarp: bool,

    /// Path to text unwarping model (uvdoc.mnn)
    #[arg(long)]
    unwarp_model: Option<PathBuf>,
    // Layout detection removed
}

#[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord, ValueEnum)]
enum OutputFormat {
    /// JSON output with full details
    Json,
    /// Plain text, one line per detected text
    Text,
    /// TSV format: text\tscore\tx1,y1,x2,y2,x3,y3,x4,y4
    Tsv,
    /// Markdown format with metadata and ordering
    Markdown,
    /// Plain text ordered by position
    TextOrdered,
}

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let cli = Cli::parse();

    // Initialize base OCR using the advanced API
    let mut ocr = RustOCR::new_ppv5(&cli.det_model, &cli.rec_model, &cli.dict)?;

    // Add optional features
    if cli.enable_orient {
        if let Some(orient_path) = cli.orient_model {
            eprintln!("Loading orientation classifier...");
            let orient_config = OrientConfig::default(orient_path);
            let orient = OrientClassifier::new(orient_config)?;
            ocr.global.use_orient = true;
            ocr = ocr.with_orient(orient);
            eprintln!("✓ Orientation classifier enabled");
        } else {
            eprintln!("Warning: --enable-orient specified but --orient-model not provided");
        }
    }

    // Layout detection removed

    // Run OCR
    eprintln!("Processing image...");
    let result = ocr.run(&cli.image)?;
    eprintln!("✓ OCR completed");

    // Output results based on format
    match cli.format {
        OutputFormat::Json => {
            let json_output = serde_json::json!({
                "boxes": result.boxes.iter().map(|b| vec![
                    vec![b[0].x, b[0].y],
                    vec![b[1].x, b[1].y],
                    vec![b[2].x, b[2].y],
                    vec![b[3].x, b[3].y],
                ]).collect::<Vec<_>>(),
                "txts": result.txts,
                "scores": result.scores,
                "orientation": result.orientation.as_ref().map(|o| o.degrees()),
                // "layout_regions" removed
                "elapse": {
                    "detection": result.elapse_det,
                    "recognition": result.elapse_rec,
                    "orientation": result.elapse_orient,
                    // "layout" removed
                },
            });
            println!("{}", serde_json::to_string_pretty(&json_output)?);
        }
        OutputFormat::Text => {
            for text in &result.txts {
                println!("{}", text);
            }
        }
        OutputFormat::Tsv => {
            for (i, text) in result.txts.iter().enumerate() {
                let bbox = &result.boxes[i];
                let score = result.scores[i];
                let box_str = format!(
                    "{:.1},{:.1},{:.1},{:.1},{:.1},{:.1},{:.1},{:.1}",
                    bbox[0].x,
                    bbox[0].y,
                    bbox[1].x,
                    bbox[1].y,
                    bbox[2].x,
                    bbox[2].y,
                    bbox[3].x,
                    bbox[3].y,
                );
                println!("{}\t{:.3}\t{}", text, score, box_str);
            }
        }
        OutputFormat::Markdown => {
            // Markdown format removed - use plain text with code block
            println!("```");
            println!("{}", result.to_raw());
            println!("```");
        }
        OutputFormat::TextOrdered => {
            println!("{}", result.to_raw());
        }
    }

    Ok(())
}