add manga-ocr

Cargo.lock

@@ -130,9 +130,9 @@ dependencies = [
 
 [[package]]
 name = "anyhow"
-version = "1.0.97"
+version = "1.0.98"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "dcfed56ad506cb2c684a14971b8861fdc3baaaae314b9e5f9bb532cbe3ba7a4f"
+checksum = "e16d2d3311acee920a9eb8d33b8cbc1787ce4a264e85f964c2404b969bdcd487"
 dependencies = [
  "backtrace",
 ]
@@ -3104,6 +3104,17 @@ version = "0.2.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "b8dd856d451cc0da70e2ef2ce95a18e39a93b7558bedf10201ad28503f918568"
 
+[[package]]
+name = "manga-ocr"
+version = "0.1.0"
+dependencies = [
+ "anyhow",
+ "clap",
+ "image",
+ "ndarray",
+ "ort",
+]
+
 [[package]]
 name = "markup5ever"
 version = "0.11.0"

Cargo.toml

@@ -1,3 +1,3 @@
 [workspace]
-members = ["src-tauri", "yolo-v8", "comic-text-detector"]
+members = ["src-tauri", "yolo-v8", "comic-text-detector", "manga-ocr"]
 resolver = "2"

manga-ocr/Cargo.toml

@@ -0,0 +1,11 @@
+[package]
+name = "manga-ocr"
+version = "0.1.0"
+edition = "2024"
+
+[dependencies]
+anyhow = "1.0.98"
+clap = { version = "4.5.36", features = ["derive"] }
+image = "0.25.6"
+ndarray = "0.16.1"
+ort = { version = "=2.0.0-rc.9", features = ["ndarray"] }

manga-ocr/src/main.rs

@@ -0,0 +1,100 @@
+use std::fs;
+
+use clap::Parser;
+use image::imageops::FilterType;
+use ndarray::{Array, s};
+use ort::inputs;
+use ort::session::Session;
+use ort::session::builder::GraphOptimizationLevel;
+
+#[derive(Parser, Debug)]
+struct Args {
+    #[arg(long)]
+    image: String,
+
+    #[arg(long)]
+    model: String,
+
+    #[arg(long)]
+    vocab: String,
+}
+
+fn main() -> anyhow::Result<()> {
+    let args = Args::parse();
+
+    let model = Session::builder()?
+        .with_optimization_level(GraphOptimizationLevel::Level3)?
+        .with_intra_threads(4)?
+        .commit_from_file(args.model)?;
+
+    let vocab = fs::read_to_string(args.vocab)
+        .map_err(|e| anyhow::anyhow!("Failed to read vocab file: {e}"))?
+        .lines()
+        .map(|s| s.to_string())
+        .collect::<Vec<_>>();
+
+    let image =
+        image::open(&args.image).map_err(|e| anyhow::anyhow!("Failed to open image: {e}"))?;
+    let image = image.grayscale().to_rgb8();
+    // Resize to 224x224
+    let image = image::imageops::resize(&image, 224, 224, FilterType::Lanczos3);
+
+    // Convert to float32 array and normalize
+    let mut tensor = Array::zeros((1, 3, 224, 224));
+    for (x, y, pixel) in image.enumerate_pixels() {
+        let x = x as usize;
+        let y = y as usize;
+
+        // Normalize from [0, 255] to [-1, 1]
+        tensor[[0, 0, y, x]] = (pixel[0] as f32 / 255.0 - 0.5) / 0.5;
+        tensor[[0, 1, y, x]] = (pixel[1] as f32 / 255.0 - 0.5) / 0.5;
+        tensor[[0, 2, y, x]] = (pixel[2] as f32 / 255.0 - 0.5) / 0.5;
+    }
+
+    // generate
+    let mut token_ids: Vec<i64> = vec![2i64]; // Start token
+
+    for _ in 0..300 {
+        // 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,
+        }?;
+
+        // Run inference
+        let outputs = model.run(inputs)?;
+
+        // Extract logits from output
+        let logits = outputs["logits"].try_extract_tensor::<f32>()?;
+
+        // Get last token logits and find argmax
+        let logits_view = logits.view();
+        let last_token_logits = logits_view.slice(s![0, -1, ..]);
+        let (token_id, _) = last_token_logits
+            .iter()
+            .enumerate()
+            .max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap())
+            .unwrap_or((0, &0.0));
+
+        token_ids.push(token_id as i64);
+
+        // Break if end token
+        if token_id as i64 == 3 {
+            break;
+        }
+    }
+
+    // decode tokens
+    let text = token_ids
+        .iter()
+        .filter(|&&id| id >= 5)
+        .filter_map(|&id| vocab.get(id as usize).cloned())
+        .collect::<Vec<_>>();
+
+    let text = text.join("");
+
+    println!("Generated text: {}", text);
+
+    Ok(())
+}

scripts/export_manga_ocr_to_onnx.py

@@ -0,0 +1,31 @@
+import torch
+from transformers import VisionEncoderDecoderModel
+
+model = VisionEncoderDecoderModel.from_pretrained("kha-white/manga-ocr-base")
+model.eval()
+
+# Dummy input for the model
+dummy_image = torch.randn(1, 3, 224, 224)
+dummy_token_ids = torch.tensor([[2]])
+
+# Export the model
+torch.onnx.export(
+    model,
+    (dummy_image, dummy_token_ids),
+    "models/manga-ocr.onnx",
+    input_names=["image", "token_ids"],
+    output_names=["logits"],
+    dynamic_axes={
+        "image": {
+            0: "batch_size",
+        },
+        "token_ids": {
+            0: "batch_size",
+            1: "sequence_length",
+        },
+        "logits": {
+            0: "batch_size",
+            1: "sequence_length",
+        },
+    },
+)

scripts/manga_ocr_onnx_inference.py

@@ -0,0 +1,97 @@
+import re
+import jaconv
+import numpy as np
+
+from onnxruntime import InferenceSession
+from PIL import Image
+
+
+class MangaOCR:
+    def __init__(self, model_path: str, vocab_path: str):
+        self.session = InferenceSession(model_path)
+        self.vocab = self._load_vocab(vocab_path)
+
+    def __call__(self, image: Image.Image) -> str:
+        image = self._preprocess(image)
+        token_ids = self._generate(image)
+        text = self._decode(token_ids)
+        text = self._postprocess(text)
+
+        return text
+
+    def _load_vocab(self, vocab_file: str) -> list[str]:
+        with open(vocab_file, "r", encoding="utf8") as f:
+            vocab = f.read().splitlines()
+
+        return vocab
+
+    def _preprocess(self, image: Image.Image) -> np.ndarray:
+        # convert to grayscale
+        image = image.convert("L").convert("RGB")
+        # resize
+        image = image.resize((224, 224), resample=2)
+        # rescale
+        image = np.array(image, dtype=np.float32)
+        image /= 255
+        # normalize
+        image = (image - 0.5) / 0.5
+        # reshape from (224, 224, 3) to (3, 224, 224)
+        image = image.transpose((2, 0, 1))
+        # add batch size
+        image = image[None]
+
+        return image
+
+    def _generate(self, image: np.ndarray) -> np.ndarray:
+        token_ids = [2]
+
+        for _ in range(300):
+            [logits] = self.session.run(
+                output_names=["logits"],
+                input_feed={
+                    "image": image,
+                    "token_ids": np.array([token_ids]),
+                },
+            )
+
+            token_id = logits[0, -1, :].argmax()
+            token_ids.append(int(token_id))
+
+            if token_id == 3:
+                break
+
+        return token_ids
+
+    def _decode(self, token_ids: list[int]) -> str:
+        text = ""
+
+        for token_id in token_ids:
+            if token_id < 5:
+                continue
+
+            text += self.vocab[token_id]
+
+        return text
+
+    def _postprocess(self, text: str) -> str:
+        text = "".join(text.split())
+        text = text.replace("…", "...")
+        text = re.sub("[・.]{2,}", lambda x: (x.end() - x.start()) * ".", text)
+        text = jaconv.h2z(text, ascii=True, digit=True)
+
+        return text
+
+
+if __name__ == "__main__":
+    import argparse
+
+    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("--vocab", type=str, help="Path to the vocabulary file")
+    args = parser.parse_args()
+
+    ocr = MangaOCR(args.model, args.vocab)
+    image = Image.open(args.image)
+    text = ocr(image)
+    print(text)