package onnx
// Ported from meikiocr (https://github.com/rtr46/meikiocr), Apache-2.0.
// Model weights: rtr46/meiki.txt.recognition.v0, LGPL-3.0.
import (
"fmt"
"image"
"math"
"sort"
"strings"
ort "github.com/yalue/onnxruntime_go"
"golang.org/x/image/draw"
)
const (
recInputHeight = 32
recInputWidth = 960
defaultRecConf = 0.1
recOverlapThresh = 0.3
)
var swappedPairs = map[string]string{
"儡傀": "傀儡",
"談冗": "冗談",
"汰淘": "淘汰",
"沱滂": "滂沱",
"攣痙": "痙攣",
"酊酩": "酩酊",
"麭麺": "麺麭",
"哭慟": "慟哭",
}
// MeikiRecognizer runs the horizontal character recognition model.
type MeikiRecognizer struct {
sess *ort.DynamicAdvancedSession
}
// NewMeikiRecognizer loads the horizontal text recognizer ONNX model.
func NewMeikiRecognizer(modelPath string) (*MeikiRecognizer, error) {
sess, err := newSession(modelPath, []string{"images", "orig_target_sizes"}, []string{"char_codes", "boxes", "scores"})
if err != nil {
return nil, fmt.Errorf("load meiki recognizer: %w", err)
}
return &MeikiRecognizer{sess: sess}, nil
}
// Close destroys the underlying ORT session.
func (r *MeikiRecognizer) Close() error {
if r.sess != nil {
err := r.sess.Destroy()
r.sess = nil
return err
}
return nil
}
type charCandidate struct {
char rune
score float32
interval [2]float64
}
// RecognizeLine recognizes one horizontal cropped text line image.
func (r *MeikiRecognizer) RecognizeLine(crop image.Image) (string, error) {
bounds := crop.Bounds()
cropW := bounds.Dx()
cropH := bounds.Dy()
if cropW <= 0 || cropH <= 0 {
return "", nil
}
scale := float64(recInputHeight) / float64(cropH)
newW := int(math.Round(float64(cropW) * scale))
newH := recInputHeight
if newW > recInputWidth {
scaleW := float64(recInputWidth) / float64(newW)
newW = recInputWidth
newH = int(math.Round(float64(recInputHeight) * scaleW))
}
if newW <= 0 || newH <= 0 {
return "", nil
}
resized := image.NewRGBA(image.Rect(0, 0, newW, newH))
draw.BiLinear.Scale(resized, resized.Bounds(), crop, bounds, draw.Over, nil)
// Prepare tensor [1, 3, 32, 960] in BGR layout
imgData := make([]float32, 1*3*recInputHeight*recInputWidth)
strideH := recInputWidth
strideC := recInputHeight * recInputWidth
for y := 0; y < newH; y++ {
for x := 0; x < newW; x++ {
c := resized.RGBAAt(x, y)
pxIdx := y*strideH + x
imgData[0*strideC+pxIdx] = float32(c.B) / 255.0
imgData[1*strideC+pxIdx] = float32(c.G) / 255.0
imgData[2*strideC+pxIdx] = float32(c.R) / 255.0
}
}
imgTensor, err := ort.NewTensor(ort.Shape{1, 3, recInputHeight, recInputWidth}, imgData)
if err != nil {
return "", fmt.Errorf("create rec images tensor: %w", err)
}
defer imgTensor.Destroy()
sizesData := []int64{recInputWidth, recInputHeight}
sizesTensor, err := ort.NewTensor(ort.Shape{1, 2}, sizesData)
if err != nil {
return "", fmt.Errorf("create rec orig_target_sizes tensor: %w", err)
}
defer sizesTensor.Destroy()
outputs := make([]ort.Value, 3)
err = r.sess.Run([]ort.Value{imgTensor, sizesTensor}, outputs)
if err != nil {
return "", fmt.Errorf("meiki rec inference: %w", err)
}
defer func() {
for _, o := range outputs {
if o != nil {
o.Destroy()
}
}
}()
charCodesTensor := outputs[0].(*ort.Tensor[int32])
boxesTensor := outputs[1].(*ort.Tensor[float32])
scoresTensor := outputs[2].(*ort.Tensor[float32])
charCodes := charCodesTensor.GetData()
boxes := boxesTensor.GetData()
scores := scoresTensor.GetData()
var candidates []charCandidate
for i, score := range scores {
if score < defaultRecConf {
continue
}
ch := rune(charCodes[i])
bOffset := i * 4
rx1 := float64(boxes[bOffset])
rx2 := float64(boxes[bOffset+2])
if rx1 >= float64(newW) {
continue
}
rx1 = math.Min(rx1, float64(newW))
rx2 = math.Min(rx2, float64(newW))
cx1 := (rx1 / float64(newW)) * float64(cropW)
cx2 := (rx2 / float64(newW)) * float64(cropW)
candidates = append(candidates, charCandidate{
char: ch,
score: score,
interval: [2]float64{cx1, cx2},
})
}
// 1D NMS: sort candidates descending by score
sort.Slice(candidates, func(i, j int) bool {
return candidates[i].score > candidates[j].score
})
var accepted []charCandidate
const epsilon = 1e-6
for _, cand := range candidates {
i1C, i2C := cand.interval[0], cand.interval[1]
lenC := i2C - i1C + epsilon
isOverlap := false
for _, acc := range accepted {
i1A, i2A := acc.interval[0], acc.interval[1]
if (i1C >= i2A) || (i1A >= i2C) {
continue
}
interStart := math.Max(i1C, i1A)
interEnd := math.Min(i2C, i2A)
interLen := math.Max(0, interEnd-interStart)
lenA := i2A - i1A + epsilon
minLen := math.Min(lenC, lenA)
if (interLen / minLen) > recOverlapThresh {
isOverlap = true
break
}
}
if !isOverlap {
accepted = append(accepted, cand)
}
}
// Sort accepted characters left-to-right along the line
sort.Slice(accepted, func(i, j int) bool {
return accepted[i].interval[0] < accepted[j].interval[0]
})
var b strings.Builder
for _, c := range accepted {
b.WriteRune(c.char)
}
text := b.String()
// Fix swapped pairs from known training corpus artifacts
for wrong, correct := range swappedPairs {
text = strings.ReplaceAll(text, wrong, correct)
}
return text, nil
}