feat: split encoder and decoder to reuse encoder outputs for manga-ocr, close #7

manga-ocr/src/main.rs

@@ -13,7 +13,10 @@ struct Args {
     image: String,
 
     #[arg(long)]
-    model: String,
+    encoder_model: String,
+
+    #[arg(long)]
+    decoder_model: String,
 
     #[arg(long)]
     vocab: String,
@@ -22,10 +25,15 @@ struct Args {
 fn main() -> anyhow::Result<()> {
     let args = Args::parse();
 
-    let model = Session::builder()?
+    let encoder_model = Session::builder()?
+        .with_optimization_level(GraphOptimizationLevel::Level3)?
+        .with_intra_threads(4)?
+        .commit_from_file(args.encoder_model)?;
+
+    let decoder_model = Session::builder()?
         .with_optimization_level(GraphOptimizationLevel::Level3)?
         .with_intra_threads(4)?
-        .commit_from_file(args.model)?;
+        .commit_from_file(args.decoder_model)?;
 
     let vocab = fs::read_to_string(args.vocab)
         .map_err(|e| anyhow::anyhow!("Failed to read vocab file: {e}"))?
@@ -51,6 +59,12 @@ fn main() -> anyhow::Result<()> {
         tensor[[0, 2, y, x]] = (pixel[2] as f32 / 255.0 - 0.5) / 0.5;
     }
 
+    let inputs = inputs! {
+        "pixel_values" => tensor.view(),
+    }?;
+    let outputs = encoder_model.run(inputs)?;
+    let encoder_hidden_states = outputs[0].try_extract_tensor::<f32>()?;
+
     // generate
     let mut token_ids: Vec<i64> = vec![2i64]; // Start token
 
@@ -58,12 +72,12 @@ fn main() -> anyhow::Result<()> {
         // Create input tensors
         let input = Array::from_shape_vec((1, token_ids.len()), token_ids.clone())?;
         let inputs = inputs! {
-            "image" => tensor.view(),
-            "token_ids" => input,
+            "encoder_hidden_states" => encoder_hidden_states.view(),
+            "input_ids" => input,
         }?;
 
         // Run inference
-        let outputs = model.run(inputs)?;
+        let outputs = decoder_model.run(inputs)?;
 
         // Extract logits from output
         let logits = outputs["logits"].try_extract_tensor::<f32>()?;

scripts/manga_ocr_onnx_inference.py

@@ -1,19 +1,27 @@
 import re
 import jaconv
 import numpy as np
+import time
 
 from onnxruntime import InferenceSession
 from PIL import Image
 
 
 class MangaOCR:
-    def __init__(self, model_path: str, vocab_path: str):
-        self.session = InferenceSession(model_path)
+    def __init__(self, encoder_model_path: str, decoder_model_path: str, vocab_path: str):
+        self.encoder_session = InferenceSession(encoder_model_path)
+        self.decoder_session = InferenceSession(decoder_model_path)
         self.vocab = self._load_vocab(vocab_path)
 
     def __call__(self, image: Image.Image) -> str:
         image = self._preprocess(image)
+
+        # count time
+        start = time.time()
         token_ids = self._generate(image)
+        end = time.time()
+        print(f"Time taken: {end - start:.2f} seconds")
+
         text = self._decode(token_ids)
         text = self._postprocess(text)
 
@@ -43,14 +51,18 @@ class MangaOCR:
         return image
 
     def _generate(self, image: np.ndarray) -> np.ndarray:
+        encoder_hidden_states = self.encoder_session.run(None, {
+            "pixel_values": image,
+        })[0]
+
         token_ids = [2]
 
         for _ in range(300):
-            [logits] = self.session.run(
-                output_names=["logits"],
-                input_feed={
-                    "image": image,
-                    "token_ids": np.array([token_ids]),
+            [logits] = self.decoder_session.run(
+                None,
+                {
+                    "encoder_hidden_states": encoder_hidden_states,
+                    "input_ids": np.array([token_ids], dtype=np.int64),
                 },
             )
 
@@ -87,11 +99,12 @@ if __name__ == "__main__":
 
     parser = argparse.ArgumentParser(description="Manga OCR with ONNX Runtime")
     parser.add_argument("--image", type=str, help="Path to the input image")
-    parser.add_argument("--model", type=str, help="Path to the ONNX model file")
+    parser.add_argument("--encoder-model", type=str, help="Path to the ONNX model file")
+    parser.add_argument("--decoder-model", type=str, help="Path to the ONNX model file")
     parser.add_argument("--vocab", type=str, help="Path to the vocabulary file")
     args = parser.parse_args()
 
-    ocr = MangaOCR(args.model, args.vocab)
+    ocr = MangaOCR(args.encoder_model, args.decoder_model, args.vocab)
     image = Image.open(args.image)
     text = ocr(image)
     print(text)

src-tauri/src/manga_ocr.rs

@@ -6,21 +6,28 @@ use ort::{inputs, session::Session};
 
 #[derive(Debug)]
 pub struct MangaOCR {
-    model: Session,
+    encoder_model: Session,
+    decoder_model: Session,
     vocab: Vec<String>,
 }
 
 impl MangaOCR {
     pub fn new() -> anyhow::Result<Self> {
         let api = Api::new()?;
-        let repo = api.model("mayocream/koharu".to_string());
-        let model_path = repo.get("manga-ocr.onnx")?;
+        let repo = api.model("mayocream/manga-ocr-onnx".to_string());
+        let encoder_model_path = repo.get("encoder_model.onnx")?;
+        let decoder_model_path = repo.get("decoder_model.onnx")?;
         let vocab_path = repo.get("vocab.txt")?;
 
-        let model = Session::builder()?
+        let encoder_model = Session::builder()?
             .with_optimization_level(ort::session::builder::GraphOptimizationLevel::Level3)?
             .with_intra_threads(thread::available_parallelism()?.get())?
-            .commit_from_file(model_path)?;
+            .commit_from_file(encoder_model_path)?;
+
+        let decoder_model = Session::builder()?
+            .with_optimization_level(ort::session::builder::GraphOptimizationLevel::Level3)?
+            .with_intra_threads(thread::available_parallelism()?.get())?
+            .commit_from_file(decoder_model_path)?;
 
         let vocab = std::fs::read_to_string(vocab_path)
             .map_err(|e| anyhow::anyhow!("Failed to read vocab file: {e}"))?
@@ -28,7 +35,11 @@ impl MangaOCR {
             .map(|s| s.to_string())
             .collect::<Vec<_>>();
 
-        Ok(Self { model, vocab })
+        Ok(Self {
+            encoder_model,
+            decoder_model,
+            vocab,
+        })
     }
 
     pub fn inference(&self, image: &image::DynamicImage) -> anyhow::Result<String> {
@@ -48,6 +59,13 @@ impl MangaOCR {
             tensor[[0, 2, y, x]] = (pixel[2] as f32 / 255.0 - 0.5) / 0.5;
         }
 
+        // save encoder hidden state
+        let inputs = inputs! {
+            "pixel_values" => tensor.view(),
+        }?;
+        let outputs = self.encoder_model.run(inputs)?;
+        let encoder_hidden_state = outputs[0].try_extract_tensor::<f32>()?;
+
         // generate
         let mut token_ids: Vec<i64> = vec![2i64]; // Start token
 
@@ -55,12 +73,12 @@ impl MangaOCR {
             // Create input tensors
             let input = ndarray::Array::from_shape_vec((1, token_ids.len()), token_ids.clone())?;
             let inputs = inputs! {
-                "image" => tensor.view(),
-                "token_ids" => input,
+                "encoder_hidden_states" => encoder_hidden_state.view(),
+                "input_ids" => input,
             }?;
 
             // Run inference
-            let outputs = self.model.run(inputs)?;
+            let outputs = self.decoder_model.run(inputs)?;
 
             // Extract logits from output
             let logits = outputs["logits"].try_extract_tensor::<f32>()?;