From 1317bb904a291eca04773ee43706a8775a098abc Mon Sep 17 00:00:00 2001
From: Phaired <65019388+Phaired@users.noreply.github.com>
Date: Fri, 20 Sep 2024 10:50:42 +0200
Subject: [PATCH] feat: cli can take custom alphabet as args (#116)

* 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>
---
 ocrs-cli/src/main.rs | 14 ++++++++++++++
 1 file changed, 14 insertions(+)

diff --git a/ocrs-cli/src/main.rs b/ocrs-cli/src/main.rs
index 96439ef..9c732df 100644
--- a/ocrs-cli/src/main.rs
+++ b/ocrs-cli/src/main.rs
@@ -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> {
@@ -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()? {
@@ -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()?);
             }
@@ -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
@@ -223,6 +235,7 @@ Advanced options:
     }
 
     Ok(Args {
+        alphabet,
         beam_search,
         debug,
         detection_model,
@@ -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 {