package onnx
// Pipeline combines meikiocr's text detector (D-FINE) with meikiocr's horizontal
// recognizer and manga-ocr's vertical reader into a single unified OCR engine.
//
// Attributions:
// - meikiocr: https://github.com/rtr46/meikiocr (Apache-2.0)
// - meikiocr weights: rtr46/meiki.text.detect.v0 & meiki.txt.recognition.v0 (LGPL-3.0)
// - manga-ocr: https://github.com/kha-white/manga-ocr (Apache-2.0)
// - manga-ocr ONNX: mayocream/manga-ocr-onnx (Apache-2.0)
import (
"fmt"
"image"
"path/filepath"
"sort"
"chirami/internal/ocr"
"golang.org/x/image/draw"
)
type Pipeline struct {
det *Detector
meiki *MeikiRecognizer
manga *MangaOCR
}
// NewPipeline initializes all models from modelDir.
func NewPipeline(modelDir string) (*Pipeline, error) {
detPath := filepath.Join(modelDir, "meiki.text.detect.v0.1.960x544.onnx")
recPath := filepath.Join(modelDir, "meiki.text.rec.v0.960x32.onnx")
encPath := filepath.Join(modelDir, "encoder_model.onnx")
decPath := filepath.Join(modelDir, "decoder_model.onnx")
vocabPath := filepath.Join(modelDir, "vocab.txt")
det, err := NewDetector(detPath)
if err != nil {
return nil, fmt.Errorf("detector: %w", err)
}
meiki, err := NewMeikiRecognizer(recPath)
if err != nil {
det.Close()
return nil, fmt.Errorf("meiki recognizer: %w", err)
}
manga, err := NewMangaOCR(encPath, decPath, vocabPath)
if err != nil {
det.Close()
meiki.Close()
return nil, fmt.Errorf("manga-ocr: %w", err)
}
return &Pipeline{
det: det,
meiki: meiki,
manga: manga,
}, nil
}
// Close releases all underlying ONNX sessions.
func (p *Pipeline) Close() error {
var errs []error
if p.det != nil {
if err := p.det.Close(); err != nil {
errs = append(errs, err)
}
p.det = nil
}
if p.meiki != nil {
if err := p.meiki.Close(); err != nil {
errs = append(errs, err)
}
p.meiki = nil
}
if p.manga != nil {
if err := p.manga.Close(); err != nil {
errs = append(errs, err)
}
p.manga = nil
}
if len(errs) > 0 {
return fmt.Errorf("close pipeline: %v", errs)
}
return nil
}
// Recognize processes an image and returns recognized text lines in reading order.
func (p *Pipeline) Recognize(img image.Image) (ocr.Result, error) {
boxes, err := p.det.Detect(img, defaultDetConf)
if err != nil {
return ocr.Result{Engine: "onnx", Error: err.Error()}, err
}
// If detector found nothing, fallback to manga-ocr on the whole image
if len(boxes) == 0 {
text, err := p.manga.Recognize(img)
if err != nil {
return ocr.Result{Engine: "onnx", Error: err.Error()}, err
}
if text == "" {
return ocr.Result{Engine: "onnx"}, nil
}
return ocr.Result{
Lines: []ocr.Line{{Text: text, Vertical: false}},
Engine: "onnx",
}, nil
}
// Filter furigana/ruby boxes so they don't end up inside the output text
filteredBoxes := FilterRubyBoxes(boxes, func(b TextBox) bool {
crop := cropImage(img, b.X1, b.Y1, b.X2, b.Y2)
var txt string
if b.IsVertical {
txt, _ = p.manga.Recognize(crop)
} else {
txt, _ = p.meiki.RecognizeLine(crop)
}
return IsKanaOnly(txt)
})
// Separate into horizontal and vertical boxes
var hBoxes []TextBox
var vBoxes []TextBox
for _, b := range filteredBoxes {
if b.IsVertical {
vBoxes = append(vBoxes, b)
} else {
hBoxes = append(hBoxes, b)
}
}
// Reading order:
// Horizontal: top-to-bottom (ascending Y1)
sort.Slice(hBoxes, func(i, j int) bool {
return hBoxes[i].Y1 < hBoxes[j].Y1
})
// Vertical: right-to-left (descending X1)
sort.Slice(vBoxes, func(i, j int) bool {
return vBoxes[i].X1 > vBoxes[j].X1
})
var lines []ocr.Line
// 1. Process horizontal lines with meiki recognizer
for _, b := range hBoxes {
crop := cropImage(img, b.X1, b.Y1, b.X2, b.Y2)
txt, err := p.meiki.RecognizeLine(crop)
if err != nil {
return ocr.Result{Engine: "onnx", Error: err.Error()}, err
}
if txt != "" {
lines = append(lines, ocr.Line{Text: txt, Vertical: false})
}
}
// 2. Process vertical columns with manga-ocr
for _, b := range vBoxes {
crop := cropImage(img, b.X1, b.Y1, b.X2, b.Y2)
txt, err := p.manga.Recognize(crop)
if err != nil {
return ocr.Result{Engine: "onnx", Error: err.Error()}, err
}
if txt != "" {
lines = append(lines, ocr.Line{Text: txt, Vertical: true})
}
}
return ocr.Result{
Lines: lines,
Engine: "onnx",
}, nil
}
func cropImage(img image.Image, x1, y1, x2, y2 int) image.Image {
type subImager interface {
SubImage(r image.Rectangle) image.Image
}
rect := image.Rect(x1, y1, x2, y2)
if si, ok := img.(subImager); ok {
return si.SubImage(rect)
}
dst := image.NewRGBA(image.Rect(0, 0, x2-x1, y2-y1))
draw.Draw(dst, dst.Bounds(), img, rect.Min, draw.Src)
return dst
}