package onnx
// Ported from meikiocr (https://github.com/rtr46/meikiocr), Apache-2.0.
// Model weights: rtr46/meiki.text.detect.v0, LGPL-3.0.
import (
"fmt"
"image"
"math"
ort "github.com/yalue/onnxruntime_go"
"golang.org/x/image/draw"
)
const (
detInputWidth = 960
detInputHeight = 544
defaultDetConf = 0.5
)
// TextBox is a detected bounding box for a line or column of text.
type TextBox struct {
X1, Y1, X2, Y2 int
Score float32
IsVertical bool
}
func (b TextBox) Width() int { return b.X2 - b.X1 }
func (b TextBox) Height() int { return b.Y2 - b.Y1 }
// Detector runs the D-FINE text detection model.
type Detector struct {
sess *ort.DynamicAdvancedSession
}
// NewDetector loads the text detector ONNX model.
func NewDetector(modelPath string) (*Detector, error) {
sess, err := newSession(modelPath, []string{"images", "orig_target_sizes"}, []string{"labels", "boxes", "scores"})
if err != nil {
return nil, fmt.Errorf("load detector session: %w", err)
}
return &Detector{sess: sess}, nil
}
// Close destroys the underlying ORT session.
func (d *Detector) Close() error {
if d.sess != nil {
err := d.sess.Destroy()
d.sess = nil
return err
}
return nil
}
// Detect finds text bounding boxes in img with confidence >= confThreshold.
func (d *Detector) Detect(img image.Image, confThreshold float32) ([]TextBox, error) {
bounds := img.Bounds()
wOrig := bounds.Dx()
hOrig := bounds.Dy()
if wOrig <= 0 || hOrig <= 0 {
return nil, nil
}
scale := math.Min(float64(detInputWidth)/float64(wOrig), float64(detInputHeight)/float64(hOrig))
wResized := int(float64(wOrig) * scale)
hResized := int(float64(hOrig) * scale)
if wResized <= 0 || hResized <= 0 {
return nil, nil
}
resized := image.NewRGBA(image.Rect(0, 0, wResized, hResized))
draw.BiLinear.Scale(resized, resized.Bounds(), img, bounds, draw.Over, nil)
// Prepare input tensor [1, 3, 544, 960] in BGR layout
imgData := make([]float32, 1*3*detInputHeight*detInputWidth)
strideH := detInputWidth
strideC := detInputHeight * detInputWidth
for y := 0; y < hResized; y++ {
for x := 0; x < wResized; x++ {
c := resized.RGBAAt(x, y)
pxIdx := y*strideH + x
// Channel 0: B, Channel 1: G, Channel 2: R
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, detInputHeight, detInputWidth}, imgData)
if err != nil {
return nil, fmt.Errorf("create images tensor: %w", err)
}
defer imgTensor.Destroy()
sizesData := []int64{int64(float64(detInputWidth) / scale), int64(float64(detInputHeight) / scale)}
sizesTensor, err := ort.NewTensor(ort.Shape{1, 2}, sizesData)
if err != nil {
return nil, fmt.Errorf("create orig_target_sizes tensor: %w", err)
}
defer sizesTensor.Destroy()
outputs := make([]ort.Value, 3)
err = d.sess.Run([]ort.Value{imgTensor, sizesTensor}, outputs)
if err != nil {
return nil, fmt.Errorf("detector inference: %w", err)
}
defer func() {
for _, o := range outputs {
if o != nil {
o.Destroy()
}
}
}()
boxesTensor := outputs[1].(*ort.Tensor[float32])
scoresTensor := outputs[2].(*ort.Tensor[float32])
boxesData := boxesTensor.GetData()
scoresData := scoresTensor.GetData()
var boxes []TextBox
for i, score := range scoresData {
if score < confThreshold {
continue
}
bOffset := i * 4
x1 := int(math.Round(float64(boxesData[bOffset])))
y1 := int(math.Round(float64(boxesData[bOffset+1])))
x2 := int(math.Round(float64(boxesData[bOffset+2])))
y2 := int(math.Round(float64(boxesData[bOffset+3])))
// Clamp to image bounds
x1 = clamp(x1, 0, wOrig)
y1 = clamp(y1, 0, hOrig)
x2 = clamp(x2, 0, wOrig)
y2 = clamp(y2, 0, hOrig)
w := x2 - x1
h := y2 - y1
if w <= 0 || h <= 0 {
continue
}
boxes = append(boxes, TextBox{
X1: x1,
Y1: y1,
X2: x2,
Y2: y2,
Score: score,
IsVertical: h > w,
})
}
return boxes, nil
}
func clamp(val, low, high int) int {
if val < low {
return low
}
if val > high {
return high
}
return val
}
// FilterRubyBoxes identifies and filters out small kana-only ruby/furigana boxes.
func FilterRubyBoxes(boxes []TextBox, isKanaOnly func(b TextBox) bool) []TextBox {
if len(boxes) <= 1 {
return boxes
}
drop := make(map[int]bool)
for i, r := range boxes {
for j, m := range boxes {
if i == j {
continue
}
// Case 1: Horizontal line with ruby above it
if !m.IsVertical && !r.IsVertical {
hM := m.Height()
hR := r.Height()
// Ruby height is significantly smaller
if float64(hR) <= float64(hM)*0.7 {
// Ruby sits immediately above main line
if r.Y2 <= m.Y1+int(float64(hM)*0.35) && r.Y1 < m.Y1 && (m.Y1-r.Y2) <= int(float64(hM)*0.85) {
oX := math.Max(0, math.Min(float64(r.X2), float64(m.X2))-math.Max(float64(r.X1), float64(m.X1)))
if oX >= float64(r.Width())*0.55 {
if isKanaOnly(r) {
drop[i] = true
break
}
}
}
}
}
// Case 2: Vertical column with ruby to the right
if m.IsVertical && r.IsVertical {
wM := m.Width()
wR := r.Width()
if float64(wR) <= float64(wM)*0.7 {
// Ruby sits to the right of vertical column
if r.X1 >= m.X1+int(float64(wM)*0.4) && (r.X1-m.X2) <= int(float64(wM)*0.85) {
oY := math.Max(0, math.Min(float64(r.Y2), float64(m.Y2))-math.Max(float64(r.Y1), float64(m.Y1)))
if oY >= float64(r.Height())*0.55 {
if isKanaOnly(r) {
drop[i] = true
break
}
}
}
}
}
}
}
var filtered []TextBox
for i, b := range boxes {
if !drop[i] {
filtered = append(filtered, b)
}
}
return filtered
}