Skip to content

Commit

Permalink
feat: cli can take custom alphabet as args (#116)
Browse files Browse the repository at this point in the history
* feat: cli can take custom alphabet as args

Add `-a, --alphabet` CLI option to specify the alphabet used by the recognition model. This is useful if the model has been trained with a custom alphabet.

---------

Co-authored-by: Robert Knight <robertknight@gmail.com>
  • Loading branch information
Phaired and robertknight authored Sep 20, 2024
1 parent 9a614c0 commit 1317bb9
Showing 1 changed file with 14 additions and 0 deletions.
14 changes: 14 additions & 0 deletions ocrs-cli/src/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -104,6 +104,10 @@ struct Args {

/// Extract each text line found and save as a PNG image.
text_line_images: bool,

/// Alphabet used by the recognition model.
/// If not provided, the default alphabet is used.
alphabet: Option<String>,
}

fn parse_args() -> Result<Args, lexopt::Error> {
Expand All @@ -119,6 +123,7 @@ fn parse_args() -> Result<Args, lexopt::Error> {
let mut text_map = false;
let mut text_mask = false;
let mut text_line_images = false;
let mut alphabet = None;

let mut parser = lexopt::Parser::from_env();
while let Some(arg) = parser.next()? {
Expand All @@ -142,6 +147,9 @@ fn parse_args() -> Result<Args, lexopt::Error> {
Short('p') | Long("png") => {
output_format = OutputFormat::Png;
}
Short('a') | Long("alphabet") => {
alphabet = Some(parser.value()?.string()?);
}
Long("rec-model") => {
recognition_model = Some(parser.value()?.string()?);
}
Expand Down Expand Up @@ -182,6 +190,10 @@ Options:
Use a custom text recognition model
-a, --alphabet \"alphabet\"
Specify the alphabet used by the recognition model
--version
Display version info
Expand Down Expand Up @@ -223,6 +235,7 @@ Advanced options:
}

Ok(Args {
alphabet,
beam_search,
debug,
detection_model,
Expand Down Expand Up @@ -278,6 +291,7 @@ fn main() -> Result<(), Box<dyn Error>> {
detection_model: Some(detection_model),
recognition_model: Some(recognition_model),
debug: args.debug,
alphabet: args.alphabet,
decode_method: if args.beam_search {
DecodeMethod::BeamSearch { width: 100 }
} else {
Expand Down

0 comments on commit 1317bb9

Please sign in to comment.