package onnx
// Ported from manga-ocr (https://github.com/kha-white/manga-ocr), Apache-2.0.
// Model weights: mayocream/manga-ocr-onnx (kha-white base ONNX), Apache-2.0.
import (
"bufio"
"fmt"
"image"
"os"
"regexp"
"strings"
ort "github.com/yalue/onnxruntime_go"
"golang.org/x/image/draw"
)
const (
mangaInputSize = 224
maxDecodeLen = 100
bosTokenID = 2 // [CLS]
eosTokenID = 3 // [SEP]
)
var dotRegex = regexp.MustCompile(`[・.]{2,}`)
// MangaOCR runs the VisionEncoderDecoder manga-ocr model.
type MangaOCR struct {
sessEnc *ort.DynamicAdvancedSession
sessDec *ort.DynamicAdvancedSession
vocab []string
}
// NewMangaOCR loads the encoder, decoder, and vocabulary.
func NewMangaOCR(encPath, decPath, vocabPath string) (*MangaOCR, error) {
vocab, err := loadVocab(vocabPath)
if err != nil {
return nil, fmt.Errorf("load vocab: %w", err)
}
sessEnc, err := newSession(encPath, []string{"pixel_values"}, []string{"last_hidden_state"})
if err != nil {
return nil, fmt.Errorf("load manga encoder: %w", err)
}
sessDec, err := newSession(decPath, []string{"input_ids", "encoder_hidden_states"}, []string{"logits"})
if err != nil {
sessEnc.Destroy()
return nil, fmt.Errorf("load manga decoder: %w", err)
}
return &MangaOCR{
sessEnc: sessEnc,
sessDec: sessDec,
vocab: vocab,
}, nil
}
func loadVocab(path string) ([]string, error) {
f, err := os.Open(path)
if err != nil {
return nil, err
}
defer f.Close()
var vocab []string
sc := bufio.NewScanner(f)
for sc.Scan() {
vocab = append(vocab, sc.Text())
}
if err := sc.Err(); err != nil {
return nil, err
}
return vocab, nil
}
// Close frees all session resources.
func (m *MangaOCR) Close() error {
var errs []error
if m.sessEnc != nil {
if err := m.sessEnc.Destroy(); err != nil {
errs = append(errs, err)
}
m.sessEnc = nil
}
if m.sessDec != nil {
if err := m.sessDec.Destroy(); err != nil {
errs = append(errs, err)
}
m.sessDec = nil
}
if len(errs) > 0 {
return fmt.Errorf("close manga-ocr: %v", errs)
}
return nil
}
// Recognize processes a single image crop or bubble.
func (m *MangaOCR) Recognize(img image.Image) (string, error) {
bounds := img.Bounds()
if bounds.Dx() <= 0 || bounds.Dy() <= 0 {
return "", nil
}
// 1. Grayscale and resize to 224x224 bilinear
resized := image.NewRGBA(image.Rect(0, 0, mangaInputSize, mangaInputSize))
draw.BiLinear.Scale(resized, resized.Bounds(), img, bounds, draw.Over, nil)
// 2. Normalize: (gray / 255.0 - 0.5) / 0.5 replicated across 3 RGB channels
pixData := make([]float32, 1*3*mangaInputSize*mangaInputSize)
channelStride := mangaInputSize * mangaInputSize
for y := 0; y < mangaInputSize; y++ {
for x := 0; x < mangaInputSize; x++ {
c := resized.RGBAAt(x, y)
// Convert RGB to grayscale (ITU-R 601-2 luma)
gray := 0.299*float32(c.R) + 0.587*float32(c.G) + 0.114*float32(c.B)
norm := (gray/255.0 - 0.5) / 0.5
idx := y*mangaInputSize + x
pixData[0*channelStride+idx] = norm
pixData[1*channelStride+idx] = norm
pixData[2*channelStride+idx] = norm
}
}
pixTensor, err := ort.NewTensor(ort.Shape{1, 3, mangaInputSize, mangaInputSize}, pixData)
if err != nil {
return "", fmt.Errorf("create pixel_values tensor: %w", err)
}
defer pixTensor.Destroy()
// 3. Encoder inference
encOut := make([]ort.Value, 1)
err = m.sessEnc.Run([]ort.Value{pixTensor}, encOut)
if err != nil {
return "", fmt.Errorf("encoder run: %w", err)
}
defer encOut[0].Destroy()
hiddenState := encOut[0].(*ort.Tensor[float32])
// 4. Greedy autoregressive decoding
inputIDs := []int64{bosTokenID}
for step := 0; step < maxDecodeLen; step++ {
seqLen := int64(len(inputIDs))
idsTensor, err := ort.NewTensor(ort.Shape{1, seqLen}, inputIDs)
if err != nil {
return "", fmt.Errorf("create input_ids tensor: %w", err)
}
decOut := make([]ort.Value, 1)
err = m.sessDec.Run([]ort.Value{idsTensor, hiddenState}, decOut)
idsTensor.Destroy()
if err != nil {
return "", fmt.Errorf("decoder run step %d: %w", step, err)
}
logitsTensor := decOut[0].(*ort.Tensor[float32])
logits := logitsTensor.GetData()
vocabSize := int(logitsTensor.GetShape()[2])
lastLogits := logits[(seqLen-1)*int64(vocabSize) : seqLen*int64(vocabSize)]
// Argmax
var bestTok int64 = 0
var bestVal float32 = -1e9
for tok, val := range lastLogits {
if val > bestVal {
bestVal = val
bestTok = int64(tok)
}
}
decOut[0].Destroy()
inputIDs = append(inputIDs, bestTok)
if bestTok == eosTokenID {
break
}
}
// 5. Detokenize and postprocess
var tokens []string
for _, id := range inputIDs {
if id >= 0 && id <= 4 { // [PAD], [UNK], [CLS], [SEP], [MASK]
continue
}
if int(id) < len(m.vocab) {
tokens = append(tokens, m.vocab[id])
}
}
text := strings.Join(tokens, "")
// Strip spaces between tokens
text = strings.Join(strings.Fields(text), "")
// Standardize ellipsis and dots
text = strings.ReplaceAll(text, "…", "...")
text = dotRegex.ReplaceAllStringFunc(text, func(m string) string {
return strings.Repeat(".", len([]rune(m)))
})
// Full-width conversion
text = H2Z(text)
return text, nil
}